[Mlir-commits] [mlir] [mlir][XeGPU][Transform] Add `xegpu.load` and `xegpu.store` coalesce patterns. (PR #200286)
Md Abdullah Shahneous Bari
llvmlistbot at llvm.org
Mon Jun 1 16:50:41 PDT 2026
https://github.com/mshahneo updated https://github.com/llvm/llvm-project/pull/200286
>From 77f431e4eceb22db844c452d8346760f732369fe Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Thu, 21 May 2026 15:26:24 +0000
Subject: [PATCH 1/5] [mlir][XeGPU][Transform] Add
`xegpu-coalesce-gather-scatter` pass.
Coalesce neighbouring lanes of xegpu.load/store into chunked accesses.
Rewrites `xegpu.load` / `xegpu.store` ops whose offsets vector describes a
contiguous-per-lane access into an equivalent op with a smaller offsets
vector and a larger `chunk_size`. The transformation reduces the number of
memory messages by having each lane fetch / write multiple contiguous
elements.
---
.../mlir/Dialect/XeGPU/Transforms/Passes.td | 37 ++
.../Dialect/XeGPU/Transforms/CMakeLists.txt | 1 +
.../Transforms/XeGPUCoalesceGatherScatter.cpp | 477 ++++++++++++++++++
.../XeGPU/coalesce-gather-scatter.mlir | 139 +++++
4 files changed, 654 insertions(+)
create mode 100644 mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
create mode 100644 mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
index 4bee1752b271e..1ac15cc80f4e7 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
@@ -118,5 +118,42 @@ def XeGPUSgToWiDistributeExperimental : Pass<"xegpu-sg-to-wi-distribute-experime
"vector::VectorDialect", "index::IndexDialect"];
}
+def XeGPUCoalesceGatherScatter : Pass<"xegpu-coalesce-gather-scatter"> {
+ let summary = "Coalesce neighbouring lanes of xegpu.load/store into chunked accesses";
+ let description = [{
+ Rewrites `xegpu.load` / `xegpu.store` ops whose offsets vector describes a
+ contiguous-per-lane access into an equivalent op with a smaller offsets
+ vector and a larger `chunk_size`. The transformation reduces the number of
+ memory messages by having each lane fetch / write multiple contiguous
+ elements.
+
+ Two index patterns are recognized:
+ - Affine offsets of the form `base + i * stride` (e.g. `vector.step` or a
+ dense constant arithmetic progression). The offsets are coalescible with
+ a factor `N` when `stride` divides `N` and the original `chunk_size`
+ multiplied by `N` does not exceed the configured maximum.
+ - All-equal offsets (e.g. `dense<0>`). All lanes load from the same
+ address; this is rewritten as a broadcast: a single load of one element
+ is materialized and the result is splatted via `vector.broadcast`.
+
+ The mask must also be uniformly `true` (a `dense<true>` constant) for the
+ coalesced lanes; partial-mask coalescing is not yet supported.
+
+ This pass is intended to run before `xegpu-propagate-layout` so that the
+ coalesced shape participates in `inst_data` selection.
+
+ Pass options:
+ - `max-chunk-size`: upper bound on the produced `chunk_size`. Defaults to
+ 8, matching typical Xe scatter chunk-size limits.
+ }];
+ let dependentDialects = [
+ "arith::ArithDialect", "memref::MemRefDialect", "xegpu::XeGPUDialect",
+ "vector::VectorDialect"];
+ let options = [Option<
+ "maxChunkSize", "max-chunk-size", "unsigned",
+ /*default=*/"8",
+ "Upper bound on the produced chunk_size when coalescing.">];
+}
+
#endif // MLIR_DIALECT_XEGPU_TRANSFORMS_PASSES_TD
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt b/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
index 0e30a6ee6e3f0..51d4eaaca6e47 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
+ XeGPUCoalesceGatherScatter.cpp
XeGPUSgToWiDistributeExperimental.cpp
XeGPUSubgroupDistribute.cpp
XeGPUUnroll.cpp
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
new file mode 100644
index 0000000000000..7a33e94cd4d84
--- /dev/null
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
@@ -0,0 +1,477 @@
+//===- XeGPUCoalesceGatherScatter.cpp - Coalesce scatter accesses --------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass coalesces neighbouring lanes of `xegpu.load` / `xegpu.store` ops
+// into accesses with a larger `chunk_size`. It targets the scatter-style
+// pointer/memref + offsets form of these ops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
+#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/Support/MathExtras.h"
+#include <optional>
+
+namespace mlir {
+namespace xegpu {
+#define GEN_PASS_DEF_XEGPUCOALESCEGATHERSCATTER
+#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"
+} // namespace xegpu
+} // namespace mlir
+
+#define DEBUG_TYPE "xegpu-coalesce-gather-scatter"
+
+using namespace mlir;
+
+namespace {
+
+/// Description of an offsets vector that has been recognized as coalescible.
+struct OffsetsPattern {
+ enum class Kind {
+ /// All lanes load the same address: rewrite as a broadcast.
+ Broadcast,
+ /// Lane `i` loads at `base + i * stride * elementSize` (in element units).
+ Affine,
+ };
+ Kind kind;
+ /// Stride in element units between adjacent lanes. Only meaningful for
+ /// `Affine`.
+ int64_t stride = 0;
+};
+
+/// Match a vector of `index` against the supported coalescible patterns.
+///
+/// Recognized:
+/// - `arith.constant dense<C>` (all lanes equal) -> Broadcast
+/// - `arith.constant dense<[a, a+s, a+2s, ...]>` -> Affine(s)
+/// - `vector.step` (lane `i` -> i) -> Affine(1)
+/// - `arith.muli %step, %splat<S>` / `arith.muli %splat<S>, %step`
+/// -> Affine(S)
+/// - `arith.addi %x, %splat<base>` / commuted -> recurse on %x
+static std::optional<OffsetsPattern> matchOffsetsPattern(Value offsets) {
+ auto vecTy = dyn_cast<VectorType>(offsets.getType());
+ if (!vecTy || vecTy.getRank() != 1)
+ return std::nullopt;
+ // Length-1 offsets vectors carry no coalescing opportunity and would let
+ // the broadcast-rewrite output re-match itself in the greedy driver.
+ if (vecTy.getNumElements() <= 1)
+ return std::nullopt;
+
+ // arith.constant dense<...>
+ if (auto cst = offsets.getDefiningOp<arith::ConstantOp>()) {
+ auto dense = dyn_cast<DenseIntElementsAttr>(cst.getValue());
+ if (!dense)
+ return std::nullopt;
+ if (dense.isSplat()) {
+ OffsetsPattern p;
+ p.kind = OffsetsPattern::Kind::Broadcast;
+ return p;
+ }
+ // Check if values form an arithmetic progression.
+ auto values = llvm::to_vector(dense.getValues<APInt>());
+ if (values.size() < 2)
+ return std::nullopt;
+ int64_t stride =
+ values[1].getSExtValue() - values[0].getSExtValue();
+ for (size_t i = 2; i < values.size(); ++i) {
+ int64_t diff = values[i].getSExtValue() - values[i - 1].getSExtValue();
+ if (diff != stride)
+ return std::nullopt;
+ }
+ OffsetsPattern p;
+ if (stride == 0) {
+ p.kind = OffsetsPattern::Kind::Broadcast;
+ } else {
+ p.kind = OffsetsPattern::Kind::Affine;
+ p.stride = stride;
+ }
+ return p;
+ }
+
+ // vector.step -> stride 1.
+ if (offsets.getDefiningOp<vector::StepOp>()) {
+ OffsetsPattern p;
+ p.kind = OffsetsPattern::Kind::Affine;
+ p.stride = 1;
+ return p;
+ }
+
+ // Helper to recognize a vector splat with constant integer value.
+ auto matchIndexSplat = [](Value v) -> std::optional<int64_t> {
+ if (auto bcast = v.getDefiningOp<vector::BroadcastOp>()) {
+ APInt c;
+ if (matchPattern(bcast.getSource(), m_ConstantInt(&c)))
+ return c.getSExtValue();
+ }
+ if (auto cst = v.getDefiningOp<arith::ConstantOp>()) {
+ if (auto dense = dyn_cast<DenseIntElementsAttr>(cst.getValue()))
+ if (dense.isSplat())
+ return dense.getSplatValue<APInt>().getSExtValue();
+ }
+ return std::nullopt;
+ };
+
+ // arith.muli with one splat operand.
+ if (auto mul = offsets.getDefiningOp<arith::MulIOp>()) {
+ Value lhs = mul.getLhs(), rhs = mul.getRhs();
+ auto lhsSplat = matchIndexSplat(lhs);
+ auto rhsSplat = matchIndexSplat(rhs);
+ Value nonSplat = lhsSplat ? rhs : (rhsSplat ? lhs : Value());
+ std::optional<int64_t> factor = lhsSplat ? lhsSplat : rhsSplat;
+ if (nonSplat && factor) {
+ auto inner = matchOffsetsPattern(nonSplat);
+ if (inner && inner->kind == OffsetsPattern::Kind::Affine) {
+ OffsetsPattern p;
+ p.kind = OffsetsPattern::Kind::Affine;
+ p.stride = inner->stride * (*factor);
+ return p;
+ }
+ if (inner && inner->kind == OffsetsPattern::Kind::Broadcast) {
+ // splat * splat is still uniform.
+ return inner;
+ }
+ }
+ }
+
+ // arith.addi with a splat operand: stride is unaffected.
+ if (auto add = offsets.getDefiningOp<arith::AddIOp>()) {
+ Value lhs = add.getLhs(), rhs = add.getRhs();
+ auto lhsSplat = matchIndexSplat(lhs);
+ auto rhsSplat = matchIndexSplat(rhs);
+ Value nonSplat = lhsSplat ? rhs : (rhsSplat ? lhs : Value());
+ if (nonSplat)
+ return matchOffsetsPattern(nonSplat);
+ }
+
+ return std::nullopt;
+}
+
+/// Returns true if `mask` is a constant `dense<true>` vector.
+static bool isAllTrueMask(Value mask) {
+ auto vecTy = dyn_cast<VectorType>(mask.getType());
+ if (!vecTy)
+ return false;
+ auto cst = mask.getDefiningOp<arith::ConstantOp>();
+ if (!cst)
+ return false;
+ auto dense = dyn_cast<DenseIntElementsAttr>(cst.getValue());
+ if (!dense || !dense.isSplat())
+ return false;
+ return dense.getSplatValue<APInt>().getBoolValue();
+}
+
+/// Compute the largest factor `N` such that:
+/// - `N` divides `numLanes`,
+/// - `N <= maxChunkSize / origChunkSize`,
+/// - lane stride (in elements) equals 1 when scaled by N (i.e. `N * stride
+/// == N * stride`, requires stride to be 1 for true contiguity at the
+/// coalesced granularity).
+///
+/// Returns std::nullopt if no useful coalescing factor exists.
+static std::optional<int64_t> chooseAffineCoalesceFactor(int64_t numLanes,
+ int64_t origChunk,
+ int64_t stride,
+ unsigned maxChunkSize) {
+ if (stride != 1)
+ return std::nullopt;
+ if (numLanes < 2)
+ return std::nullopt;
+ if (origChunk < 1)
+ origChunk = 1;
+ int64_t budget = static_cast<int64_t>(maxChunkSize) / origChunk;
+ if (budget < 2)
+ return std::nullopt;
+ // Pick the largest power-of-two factor that divides numLanes and fits
+ // budget.
+ int64_t factor = 1;
+ for (int64_t f = std::min<int64_t>(budget, numLanes); f >= 2; f /= 2) {
+ if (numLanes % f == 0) {
+ factor = f;
+ break;
+ }
+ }
+ if (factor < 2)
+ return std::nullopt;
+ return factor;
+}
+
+/// Pick the new offsets value: a sub-vector of length `newLen` extracted from
+/// the original offsets, taking every `factor`-th element. For affine offsets
+/// generated from `vector.step` + scalar adds/muls, the original offsets
+/// already encode the right values; the simplest valid construction is to
+/// rebuild them using the same defining chain at the smaller length.
+///
+/// We take a pragmatic approach: emit
+/// %step = vector.step : vector<newLen x index>
+/// %scaled = arith.muli %step, %splat<factor*stride_element>
+/// %final = arith.addi %scaled, %splat<lane0_offset>
+///
+/// where `lane0_offset` is `offsets[0]` extracted at runtime via
+/// `vector.extract`. For dense-constant offsets we just emit a smaller dense
+/// constant directly.
+static Value buildCoalescedAffineOffsets(OpBuilder &b, Location loc,
+ Value origOffsets, int64_t newLen,
+ int64_t factor) {
+ auto idxTy = b.getIndexType();
+ auto newVecTy = VectorType::get({newLen}, idxTy);
+
+ // Fast path: dense constant offsets -> emit a smaller dense constant.
+ if (auto cst = origOffsets.getDefiningOp<arith::ConstantOp>()) {
+ if (auto dense = dyn_cast<DenseIntElementsAttr>(cst.getValue())) {
+ auto srcVals = llvm::to_vector(dense.getValues<APInt>());
+ SmallVector<APInt> newVals;
+ newVals.reserve(newLen);
+ for (int64_t i = 0; i < newLen; ++i)
+ newVals.push_back(srcVals[i * factor]);
+ auto newAttr =
+ DenseIntElementsAttr::get(newVecTy, ArrayRef<APInt>(newVals));
+ return arith::ConstantOp::create(b, loc, newAttr);
+ }
+ }
+
+ // General path: rebuild using vector.step + scalar broadcast/muli/addi.
+ // Lane-0 offset is the first element of the original offsets.
+ Value zero = arith::ConstantIndexOp::create(b, loc, 0);
+ Value lane0 =
+ vector::ExtractOp::create(b, loc, origOffsets, ArrayRef<int64_t>{0});
+ Value step = vector::StepOp::create(b, loc, newVecTy);
+ Value factorSplat = arith::ConstantOp::create(
+ b, loc,
+ DenseIntElementsAttr::get(newVecTy,
+ APInt(64, factor, /*isSigned=*/true)));
+ // Note: vector.step yields index, factorSplat must also be index-typed.
+ // Build factor splat as index using broadcast of scalar.
+ factorSplat = vector::BroadcastOp::create(
+ b, loc, newVecTy,
+ arith::ConstantIndexOp::create(b, loc, factor).getResult());
+ Value scaled = arith::MulIOp::create(b, loc, step, factorSplat);
+ Value baseSplat = vector::BroadcastOp::create(b, loc, newVecTy, lane0);
+ Value result = arith::AddIOp::create(b, loc, scaled, baseSplat);
+ (void)zero;
+ return result;
+}
+
+/// Build a `dense<true>` mask vector of length `newLen`.
+static Value buildAllTrueMask(OpBuilder &b, Location loc, int64_t newLen) {
+ auto vecTy = VectorType::get({newLen}, b.getI1Type());
+ auto attr = DenseIntElementsAttr::get(vecTy, true);
+ return arith::ConstantOp::create(b, loc, attr);
+}
+
+/// Pattern that coalesces a `xegpu.load` op.
+struct CoalesceLoadPattern final : OpRewritePattern<xegpu::LoadGatherOp> {
+ CoalesceLoadPattern(MLIRContext *ctx, unsigned maxChunkSize)
+ : OpRewritePattern(ctx), maxChunkSize(maxChunkSize) {}
+
+ LogicalResult matchAndRewrite(xegpu::LoadGatherOp op,
+ PatternRewriter &rewriter) const override {
+ // Only the source-as-pointer/memref form has a vector offsets operand we
+ // can analyze; the tensor_desc form has no offsets here.
+ auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
+ if (!offsetsTy || offsetsTy.getRank() != 1)
+ return rewriter.notifyMatchFailure(op, "expected 1-D vector offsets");
+ if (offsetsTy.getNumElements() <= 1)
+ return rewriter.notifyMatchFailure(op, "nothing to coalesce");
+ auto valueTy = op.getValueType();
+ if (!valueTy || valueTy.getRank() != 1)
+ return rewriter.notifyMatchFailure(op, "expected 1-D vector value");
+ if (!isAllTrueMask(op.getMask()))
+ return rewriter.notifyMatchFailure(op, "non-uniform mask");
+
+ auto pattern = matchOffsetsPattern(op.getOffsets());
+ if (!pattern)
+ return rewriter.notifyMatchFailure(op, "offsets not coalescible");
+
+ int64_t numLanes = offsetsTy.getNumElements();
+ int64_t origChunk =
+ static_cast<int64_t>(op.getChunkSize().value_or(1));
+
+ Location loc = op.getLoc();
+
+ if (pattern->kind == OffsetsPattern::Kind::Broadcast) {
+ // All lanes load the same address: emit a single scalar load and
+ // broadcast.
+ Value scalarOffset = vector::ExtractOp::create(
+ rewriter, loc, op.getOffsets(), ArrayRef<int64_t>{0});
+ Value scalarMask = arith::ConstantOp::create(
+ rewriter, loc, rewriter.getOneAttr(rewriter.getI1Type()));
+ // Result of the scalar load matches the original element type with the
+ // chunk dimension preserved if any (1-D value, length = chunk_size or 1).
+ int64_t scalarLen = numLanes; // value length per lane is implicit; for
+ // 1-D the whole vector represents lanes.
+ // Build a length-1 offsets vector to keep types consistent (the op
+ // accepts scalar offsets too, but our incoming form is vector).
+ auto idxVecTy = VectorType::get({1}, rewriter.getIndexType());
+ auto maskVecTy = VectorType::get({1}, rewriter.getI1Type());
+ Value newOffsets = vector::BroadcastOp::create(
+ rewriter, loc, idxVecTy, scalarOffset);
+ Value newMask = arith::ConstantOp::create(
+ rewriter, loc, DenseIntElementsAttr::get(maskVecTy, true));
+ auto newValueTy =
+ VectorType::get({1}, valueTy.getElementType());
+ auto newLoad = xegpu::LoadGatherOp::create(
+ rewriter, loc, newValueTy, op.getSource(), newOffsets, newMask,
+ /*chunk_size=*/IntegerAttr(), op.getL1HintAttr(),
+ op.getL2HintAttr(), op.getL3HintAttr(),
+ /*layout=*/xegpu::DistributeLayoutAttr());
+ // Splat to original shape.
+ Value scalar = vector::ExtractOp::create(
+ rewriter, loc, newLoad.getResult(), ArrayRef<int64_t>{0});
+ Value bcast =
+ vector::BroadcastOp::create(rewriter, loc, valueTy, scalar);
+ rewriter.replaceOp(op, bcast);
+ (void)scalarMask;
+ (void)scalarLen;
+ return success();
+ }
+
+ // Affine path.
+ auto factorOpt = chooseAffineCoalesceFactor(numLanes, origChunk,
+ pattern->stride, maxChunkSize);
+ if (!factorOpt)
+ return rewriter.notifyMatchFailure(op, "no useful coalesce factor");
+ int64_t factor = *factorOpt;
+ int64_t newLanes = numLanes / factor;
+ int64_t newChunk = origChunk * factor;
+
+ Value newOffsets = buildCoalescedAffineOffsets(rewriter, loc,
+ op.getOffsets(),
+ newLanes, factor);
+ Value newMask = buildAllTrueMask(rewriter, loc, newLanes);
+
+ // The verifier requires `chunk_size > 1` to be paired with a 2-D value
+ // vector of shape `<lanes x chunk>`. Build that 2-D type and shape_cast
+ // back to the original 1-D shape for the consumer.
+ SmallVector<int64_t> newShape;
+ if (valueTy.getRank() == 1) {
+ newShape = {newLanes, newChunk};
+ } else {
+ newShape = llvm::to_vector(valueTy.getShape());
+ newShape[0] = newLanes;
+ newShape.back() = newChunk;
+ }
+ auto newValueTy = VectorType::get(newShape, valueTy.getElementType());
+
+ auto newChunkAttr = rewriter.getI64IntegerAttr(newChunk);
+ auto newLoad = xegpu::LoadGatherOp::create(
+ rewriter, loc, newValueTy, op.getSource(), newOffsets, newMask,
+ newChunkAttr, op.getL1HintAttr(), op.getL2HintAttr(),
+ op.getL3HintAttr(), /*layout=*/xegpu::DistributeLayoutAttr());
+
+ if (newValueTy == valueTy) {
+ rewriter.replaceOp(op, newLoad.getResult());
+ } else {
+ Value reshaped = vector::ShapeCastOp::create(
+ rewriter, loc, valueTy, newLoad.getResult());
+ rewriter.replaceOp(op, reshaped);
+ }
+ return success();
+ }
+
+ unsigned maxChunkSize;
+};
+
+/// Pattern that coalesces a `xegpu.store` op.
+struct CoalesceStorePattern final : OpRewritePattern<xegpu::StoreScatterOp> {
+ CoalesceStorePattern(MLIRContext *ctx, unsigned maxChunkSize)
+ : OpRewritePattern(ctx), maxChunkSize(maxChunkSize) {}
+
+ LogicalResult matchAndRewrite(xegpu::StoreScatterOp op,
+ PatternRewriter &rewriter) const override {
+ auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
+ if (!offsetsTy || offsetsTy.getRank() != 1)
+ return rewriter.notifyMatchFailure(op, "expected 1-D vector offsets");
+ if (offsetsTy.getNumElements() <= 1)
+ return rewriter.notifyMatchFailure(op, "nothing to coalesce");
+ auto valueTy = op.getValueType();
+ if (!valueTy || valueTy.getRank() != 1)
+ return rewriter.notifyMatchFailure(op, "expected 1-D vector value");
+ if (!isAllTrueMask(op.getMask()))
+ return rewriter.notifyMatchFailure(op, "non-uniform mask");
+
+ auto pattern = matchOffsetsPattern(op.getOffsets());
+ if (!pattern)
+ return rewriter.notifyMatchFailure(op, "offsets not coalescible");
+
+ // For the broadcast case, all lanes write the same address. That is not a
+ // meaningful coalesce target (last writer wins, semantics-preserving
+ // rewrite would need a single arbitrary lane to win). Skip.
+ if (pattern->kind == OffsetsPattern::Kind::Broadcast)
+ return rewriter.notifyMatchFailure(
+ op, "all-equal offsets on store would be ambiguous");
+
+ int64_t numLanes = offsetsTy.getNumElements();
+ int64_t origChunk =
+ static_cast<int64_t>(op.getChunkSize().value_or(1));
+ auto factorOpt = chooseAffineCoalesceFactor(numLanes, origChunk,
+ pattern->stride, maxChunkSize);
+ if (!factorOpt)
+ return rewriter.notifyMatchFailure(op, "no useful coalesce factor");
+ int64_t factor = *factorOpt;
+ int64_t newLanes = numLanes / factor;
+ int64_t newChunk = origChunk * factor;
+
+ Location loc = op.getLoc();
+ Value newOffsets = buildCoalescedAffineOffsets(rewriter, loc,
+ op.getOffsets(),
+ newLanes, factor);
+ Value newMask = buildAllTrueMask(rewriter, loc, newLanes);
+
+ auto newChunkAttr = rewriter.getI64IntegerAttr(newChunk);
+ // The verifier requires `chunk_size > 1` to be paired with a 2-D value
+ // vector of shape `<lanes x chunk>`; shape_cast the original 1-D value
+ // into that shape.
+ SmallVector<int64_t> newShape;
+ if (valueTy.getRank() == 1)
+ newShape = {newLanes, newChunk};
+ else {
+ newShape = llvm::to_vector(valueTy.getShape());
+ newShape[0] = newLanes;
+ newShape.back() = newChunk;
+ }
+ auto newValTy = VectorType::get(newShape, valueTy.getElementType());
+ Value newValue = op.getValue();
+ if (newValTy != valueTy)
+ newValue =
+ vector::ShapeCastOp::create(rewriter, loc, newValTy, op.getValue());
+ xegpu::StoreScatterOp::create(rewriter, loc, newValue, op.getDest(),
+ newOffsets, newMask, newChunkAttr,
+ op.getL1HintAttr(), op.getL2HintAttr(),
+ op.getL3HintAttr(),
+ /*layout=*/xegpu::DistributeLayoutAttr());
+ rewriter.eraseOp(op);
+ return success();
+ }
+
+ unsigned maxChunkSize;
+};
+
+struct XeGPUCoalesceGatherScatterPass final
+ : public xegpu::impl::XeGPUCoalesceGatherScatterBase<
+ XeGPUCoalesceGatherScatterPass> {
+ using XeGPUCoalesceGatherScatterBase::XeGPUCoalesceGatherScatterBase;
+
+ void runOnOperation() override {
+ MLIRContext *ctx = &getContext();
+ RewritePatternSet patterns(ctx);
+ patterns.add<CoalesceLoadPattern, CoalesceStorePattern>(ctx, maxChunkSize);
+ if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+ return signalPassFailure();
+ }
+};
+
+} // namespace
diff --git a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
new file mode 100644
index 0000000000000..be205f152604e
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
@@ -0,0 +1,139 @@
+// RUN: mlir-opt -split-input-file -xegpu-coalesce-gather-scatter %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -xegpu-coalesce-gather-scatter="max-chunk-size=4" %s | FileCheck --check-prefix=CHECK4 %s
+
+// -----
+// vector.step offsets -> stride 1, fully coalescible.
+// 32 lanes, max-chunk-size = 8 -> factor 8 -> 4 lanes, chunk_size = 8.
+// CHECK-LABEL: @load_step_offsets
+// CHECK-SAME: (%[[ARG:.*]]: i64) -> vector<32xf32>
+// CHECK: %[[STEP:.*]] = vector.step : vector<4xindex>
+// CHECK: %[[LOAD:.*]] = xegpu.load %[[ARG]][{{.*}}], {{.*}} <{chunk_size = 8 : i64}>
+// CHECK-SAME: : i64, vector<4xindex>, vector<4xi1> -> vector<4x8xf32>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] : vector<4x8xf32> to vector<32xf32>
+// CHECK: return %[[CAST]]
+func.func @load_step_offsets(%ptr: i64) -> vector<32xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// max-chunk-size = 4: factor 4 -> 8 lanes, chunk_size = 4.
+// CHECK4-LABEL: @load_step_offsets_chunk4
+// CHECK4: xegpu.load {{.*}} <{chunk_size = 4 : i64}>
+// CHECK4-SAME: : i64, vector<8xindex>, vector<8xi1> -> vector<8x4xf32>
+// CHECK4: vector.shape_cast {{.*}} : vector<8x4xf32> to vector<32xf32>
+func.func @load_step_offsets_chunk4(%ptr: i64) -> vector<32xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// Dense constant arithmetic progression with stride 1.
+// CHECK-LABEL: @load_dense_ap_offsets
+// CHECK: %[[CST:.*]] = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
+// CHECK: xegpu.load {{.*}}[%[[CST]]], {{.*}} <{chunk_size = 8 : i64}>
+// CHECK-SAME: : i64, vector<4xindex>, vector<4xi1> -> vector<4x8xi32>
+// CHECK: vector.shape_cast {{.*}} : vector<4x8xi32> to vector<32xi32>
+func.func @load_dense_ap_offsets(%ptr: i64) -> vector<32xi32> {
+ %offsets = arith.constant dense<[
+ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
+ 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31
+ ]> : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xi32>
+ return %v : vector<32xi32>
+}
+
+// -----
+// All-equal offsets -> broadcast load (single load + vector.broadcast).
+// CHECK-LABEL: @load_broadcast_offsets
+// CHECK: %[[L:.*]] = xegpu.load
+// CHECK-SAME: : i64, vector<1xindex>, vector<1xi1> -> vector<1xf32>
+// CHECK: %[[E:.*]] = vector.extract %[[L]][0]
+// CHECK: %[[B:.*]] = vector.broadcast %[[E]] : f32 to vector<32xf32>
+// CHECK: return %[[B]]
+func.func @load_broadcast_offsets(%ptr: i64) -> vector<32xf32> {
+ %offsets = arith.constant dense<0> : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// Non-stride-1 (stride 4) offsets: not contiguous, not coalesced.
+// CHECK-LABEL: @load_stride4_unchanged
+// CHECK: xegpu.load
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+func.func @load_stride4_unchanged(%ptr: i64) -> vector<32xf32> {
+ %c4 = arith.constant 4 : index
+ %step = vector.step : vector<32xindex>
+ %splat = vector.broadcast %c4 : index to vector<32xindex>
+ %offsets = arith.muli %step, %splat : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// Non-uniform mask: not coalesced.
+// CHECK-LABEL: @load_partial_mask_unchanged
+// CHECK: xegpu.load
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+func.func @load_partial_mask_unchanged(%ptr: i64, %mask: vector<32xi1>) -> vector<32xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// Store with vector.step offsets coalesces.
+// CHECK-LABEL: @store_step_offsets
+// CHECK: vector.shape_cast {{.*}} : vector<32xf32> to vector<4x8xf32>
+// CHECK: xegpu.store
+// CHECK-SAME: <{chunk_size = 8 : i64}>
+// CHECK-SAME: vector<4x8xf32>, i64, vector<4xindex>, vector<4xi1>
+func.func @store_step_offsets(%ptr: i64, %v: vector<32xf32>) {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ xegpu.store %v, %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
+ return
+}
+
+// -----
+// Store with all-equal offsets is left alone (ambiguous semantics).
+// CHECK-LABEL: @store_broadcast_offsets_unchanged
+// CHECK: xegpu.store
+// CHECK-SAME: vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
+func.func @store_broadcast_offsets_unchanged(%ptr: i64, %v: vector<32xf32>) {
+ %offsets = arith.constant dense<0> : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ xegpu.store %v, %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
+ return
+}
+
+// -----
+// memref-source variant of load coalesces too.
+// CHECK-LABEL: @load_memref_step
+// CHECK: xegpu.load
+// CHECK-SAME: <{chunk_size = 8 : i64}>
+// CHECK-SAME: : memref<1024xf32>, vector<4xindex>, vector<4xi1> -> vector<4x8xf32>
+// CHECK: vector.shape_cast {{.*}} : vector<4x8xf32> to vector<32xf32>
+func.func @load_memref_step(%m: memref<1024xf32>) -> vector<32xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %m[%offsets], %mask <{chunk_size = 1 : i64}>
+ : memref<1024xf32>, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
>From 7cf101969884b7d8bea14f0f6db77dbfc3bbfe3e Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Tue, 26 May 2026 18:32:57 +0000
Subject: [PATCH 2/5] [mlir][XeGPU][Transform] Use AxisInfo dataflow +
lane_data layout for coalescing.
Rework `xegpu-coalesce-gather-scatter` along two axes:
1. Replace ad-hoc offsets-producer pattern matching with a small
XeGPU-local sparse-forward dataflow analysis modeled on Triton's
`AxisInfo` (per-dim contiguity / constancy / divisibility /
knownConstant). Visitors cover `vector.step`, dense constants
(including N-D dense arithmetic-progression detection along the
innermost dim), `vector.broadcast`, `vector.shape_cast` (identity-like
and general linear reshape), `arith.addi`/`subi`/`muli`, and
`arith.index_cast`/`index_castui`.
2. Stop rewriting the op shape (`chunk_size` + `vector.shape_cast`).
Instead, attach a `xegpu::LayoutAttr` with `lane_data[innermost] =
factor` and `lane_layout[innermost] = innerLanes/factor` on the
original op, leaving its value type, `chunk_size`, offsets, and mask
unchanged. The downstream WG-to-SG / SG-to-Lane distribution passes
already consume `lane_data` to issue chunked memory messages.
The analysis works on N-D vectors. The coalescing decision is computed
against the innermost dim of the offsets vector. New lit tests cover:
2-D leading-unit-dim, true 2-D dense-AP, 1-D step reshaped to 2-D, an
N-D non-AP negative case, and an N-D store variant.
The all-equal-offsets broadcast-load case still rewrites to a length-1
`xegpu.load` + `vector.broadcast` (there is no layout-only encoding for
"all lanes load the same scalar"). Stores with all-equal offsets are
still skipped.
Co-Authored-By: Claude Opus 4.7 <noreply at anthropic.com>
---
.../Transforms/XeGPUCoalesceGatherScatter.cpp | 978 ++++++++++++------
.../XeGPU/coalesce-gather-scatter.mlir | 145 ++-
2 files changed, 750 insertions(+), 373 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
index 7a33e94cd4d84..91808b711cf12 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
@@ -7,11 +7,31 @@
//===----------------------------------------------------------------------===//
//
// This pass coalesces neighbouring lanes of `xegpu.load` / `xegpu.store` ops
-// into accesses with a larger `chunk_size`. It targets the scatter-style
-// pointer/memref + offsets form of these ops.
+// so that each lane handles `N` contiguous elements along the innermost
+// dimension. The decision is driven by a small XeGPU-local axis-info
+// dataflow analysis modeled on Triton's `AxisInfo` (`contiguity`,
+// `constancy`, `divisibility`) and is applied by attaching a
+// `lane_data` layout to the original op. The actual memory-message rewrite
+// is left to the downstream WG-to-SG / SG-to-Lane distribution passes,
+// which interpret `lane_data`.
+//
+// The analysis tracks per-axis information for vectors of integer / index
+// type at any rank. The coalescing decision is computed against the
+// innermost dimension. 2-D offsets vectors with a leading unit dimension
+// (e.g. `vector<1x32xindex>`) are handled by treating the inner dim as the
+// lane dim.
+//
+// The `Broadcast` case (constancy along the innermost dim equals the inner
+// length) is special: there is no layout-only encoding for "all lanes load
+// the same scalar", so for loads we still rewrite to a length-1
+// `xegpu.load` followed by `vector.broadcast`. Stores in this shape are
+// skipped (last-writer-wins is ambiguous).
//
//===----------------------------------------------------------------------===//
+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
+#include "mlir/Analysis/DataFlowFramework.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
@@ -23,7 +43,9 @@
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/bit.h"
#include "llvm/Support/MathExtras.h"
+#include <numeric>
#include <optional>
namespace mlir {
@@ -37,127 +59,525 @@ namespace xegpu {
using namespace mlir;
-namespace {
+// AxisInfo and AxisInfoAnalysis are intentionally placed in a named namespace
+// (not anonymous) so the `dataflow::Lattice<AxisInfo>` template instantiation
+// gets a stable, externally-visible name. The TypeID machinery requires that
+// for `MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID` to apply.
+namespace mlir::xegpu::detail::axis_dataflow {
-/// Description of an offsets vector that has been recognized as coalescible.
-struct OffsetsPattern {
- enum class Kind {
- /// All lanes load the same address: rewrite as a broadcast.
- Broadcast,
- /// Lane `i` loads at `base + i * stride * elementSize` (in element units).
- Affine,
- };
- Kind kind;
- /// Stride in element units between adjacent lanes. Only meaningful for
- /// `Affine`.
- int64_t stride = 0;
-};
+//===----------------------------------------------------------------------===//
+// AxisInfo: per-axis contiguity / constancy / divisibility lattice.
+//===----------------------------------------------------------------------===//
-/// Match a vector of `index` against the supported coalescible patterns.
+/// Sentinel "very large" value for unconstrained dimensions. Any real shape
+/// is far smaller, so component-wise `min` will collapse this to the truth.
+static constexpr int64_t kAxisInfoTop = 1LL << 30;
+
+/// Per-dimension axis information for an SSA value of integer / index type.
+/// - `contiguity[d]`: the largest N such that consecutive lanes along
+/// dimension `d` differ by exactly 1 in runs of length `N` (lane-stride
+/// 1 contiguity).
+/// - `constancy[d]`: the largest N such that consecutive lanes along
+/// dimension `d` are all equal in runs of length `N`.
+/// - `divisibility[d]`: a power-of-two divisor of every element along
+/// dimension `d`.
+/// - `knownConstant`: scalar value if the entire vector is uniformly
+/// known to be a single constant.
///
-/// Recognized:
-/// - `arith.constant dense<C>` (all lanes equal) -> Broadcast
-/// - `arith.constant dense<[a, a+s, a+2s, ...]>` -> Affine(s)
-/// - `vector.step` (lane `i` -> i) -> Affine(1)
-/// - `arith.muli %step, %splat<S>` / `arith.muli %splat<S>, %step`
-/// -> Affine(S)
-/// - `arith.addi %x, %splat<base>` / commuted -> recurse on %x
-static std::optional<OffsetsPattern> matchOffsetsPattern(Value offsets) {
- auto vecTy = dyn_cast<VectorType>(offsets.getType());
- if (!vecTy || vecTy.getRank() != 1)
- return std::nullopt;
- // Length-1 offsets vectors carry no coalescing opportunity and would let
- // the broadcast-rewrite output re-match itself in the greedy driver.
- if (vecTy.getNumElements() <= 1)
- return std::nullopt;
-
- // arith.constant dense<...>
- if (auto cst = offsets.getDefiningOp<arith::ConstantOp>()) {
- auto dense = dyn_cast<DenseIntElementsAttr>(cst.getValue());
- if (!dense)
- return std::nullopt;
+/// Pessimistic / entry value: contiguity=1, constancy=1, divisibility=1.
+struct AxisInfo {
+ SmallVector<int64_t> contiguity;
+ SmallVector<int64_t> constancy;
+ SmallVector<int64_t> divisibility;
+ std::optional<int64_t> knownConstant;
+
+ AxisInfo() = default;
+
+ static AxisInfo getPessimistic(unsigned rank) {
+ AxisInfo v;
+ v.contiguity.assign(rank, 1);
+ v.constancy.assign(rank, 1);
+ v.divisibility.assign(rank, 1);
+ return v;
+ }
+
+ unsigned getRank() const { return contiguity.size(); }
+ bool isInitialized() const { return getRank() > 0; }
+
+ bool operator==(const AxisInfo &rhs) const {
+ return contiguity == rhs.contiguity && constancy == rhs.constancy &&
+ divisibility == rhs.divisibility &&
+ knownConstant == rhs.knownConstant;
+ }
+
+ /// Conservative join. Two values reaching the same SSA value via different
+ /// control-flow paths must agree on what holds.
+ static AxisInfo join(const AxisInfo &lhs, const AxisInfo &rhs) {
+ if (!lhs.isInitialized())
+ return rhs;
+ if (!rhs.isInitialized())
+ return lhs;
+ assert(lhs.getRank() == rhs.getRank());
+ AxisInfo out;
+ unsigned r = lhs.getRank();
+ out.contiguity.resize(r);
+ out.constancy.resize(r);
+ out.divisibility.resize(r);
+ for (unsigned d = 0; d < r; ++d) {
+ out.contiguity[d] = std::min(lhs.contiguity[d], rhs.contiguity[d]);
+ out.constancy[d] = std::min(lhs.constancy[d], rhs.constancy[d]);
+ out.divisibility[d] = std::gcd(lhs.divisibility[d], rhs.divisibility[d]);
+ }
+ if (lhs.knownConstant && rhs.knownConstant &&
+ *lhs.knownConstant == *rhs.knownConstant)
+ out.knownConstant = lhs.knownConstant;
+ return out;
+ }
+
+ void print(raw_ostream &os) const {
+ os << "contiguity=[";
+ llvm::interleaveComma(contiguity, os);
+ os << "] constancy=[";
+ llvm::interleaveComma(constancy, os);
+ os << "] divisibility=[";
+ llvm::interleaveComma(divisibility, os);
+ os << "]";
+ if (knownConstant)
+ os << " const=" << *knownConstant;
+ }
+};
+
+using AxisInfoLattice = dataflow::Lattice<AxisInfo>;
+
+/// Power-of-two divisor of `v`. Returns `kAxisInfoTop` when `v == 0`.
+static int64_t highestPow2Divisor(int64_t v) {
+ if (v == 0)
+ return kAxisInfoTop;
+ uint64_t u = static_cast<uint64_t>(std::abs(v));
+ return static_cast<int64_t>(u & (~u + 1));
+}
+
+/// Initial lattice value for an SSA value when no transfer function applies.
+static AxisInfo entryStateFor(Value v) {
+ if (auto vt = dyn_cast<VectorType>(v.getType()))
+ return AxisInfo::getPessimistic(vt.getRank());
+ return AxisInfo::getPessimistic(1);
+}
+
+/// AxisInfo for a tensor that is constant `c` everywhere.
+static AxisInfo splatAxisInfo(ArrayRef<int64_t> shape, int64_t c) {
+ AxisInfo v;
+ unsigned r = shape.size();
+ v.contiguity.assign(r, 1);
+ v.constancy.assign(shape.begin(), shape.end());
+ v.divisibility.assign(r, highestPow2Divisor(c));
+ v.knownConstant = c;
+ return v;
+}
+
+/// Sparse forward dataflow analysis that computes `AxisInfo` for vector
+/// values reachable from the entry of the analyzed op.
+class AxisInfoAnalysis
+ : public dataflow::SparseForwardDataFlowAnalysis<AxisInfoLattice> {
+public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AxisInfoAnalysis)
+ using SparseForwardDataFlowAnalysis::SparseForwardDataFlowAnalysis;
+
+ LogicalResult visitOperation(Operation *op,
+ ArrayRef<const AxisInfoLattice *> operands,
+ ArrayRef<AxisInfoLattice *> results) override {
+ if (auto step = dyn_cast<vector::StepOp>(op))
+ return visitStep(step, results);
+ if (auto cst = dyn_cast<arith::ConstantOp>(op))
+ return visitConstant(cst, results);
+ if (auto bcast = dyn_cast<vector::BroadcastOp>(op))
+ return visitBroadcast(bcast, operands, results);
+ if (auto sc = dyn_cast<vector::ShapeCastOp>(op))
+ return visitShapeCast(sc, operands, results);
+ if (auto add = dyn_cast<arith::AddIOp>(op))
+ return visitAddSub</*IsSub=*/false>(add, operands, results);
+ if (auto sub = dyn_cast<arith::SubIOp>(op))
+ return visitAddSub</*IsSub=*/true>(sub, operands, results);
+ if (auto mul = dyn_cast<arith::MulIOp>(op))
+ return visitMul(mul, operands, results);
+ if (auto cast = dyn_cast<arith::IndexCastOp>(op))
+ return visitPassThrough(cast, operands, results);
+ if (auto cast = dyn_cast<arith::IndexCastUIOp>(op))
+ return visitPassThrough(cast, operands, results);
+ setAllPessimistic(op, results);
+ return success();
+ }
+
+ void setToEntryState(AxisInfoLattice *lattice) override {
+ propagateIfChanged(lattice,
+ lattice->join(entryStateFor(lattice->getAnchor())));
+ }
+
+private:
+ void setAllPessimistic(Operation *op, ArrayRef<AxisInfoLattice *> results) {
+ for (auto [r, lat] : llvm::zip(op->getResults(), results)) {
+ AxisInfo state = entryStateFor(r);
+ propagateIfChanged(lat, lat->join(state));
+ }
+ }
+
+ // vector.step is always 1-D and produces [0, 1, ..., n-1].
+ LogicalResult visitStep(vector::StepOp op,
+ ArrayRef<AxisInfoLattice *> results) {
+ auto vt = cast<VectorType>(op.getType());
+ int64_t n = vt.getNumElements();
+ AxisInfo v;
+ v.contiguity = {n};
+ v.constancy = {1};
+ v.divisibility = {kAxisInfoTop};
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+
+ LogicalResult visitConstant(arith::ConstantOp op,
+ ArrayRef<AxisInfoLattice *> results) {
+ auto vt = dyn_cast<VectorType>(op.getType());
+ if (!vt) {
+ if (auto intAttr = dyn_cast<IntegerAttr>(op.getValue())) {
+ int64_t c = intAttr.getValue().getSExtValue();
+ AxisInfo v;
+ v.contiguity = {1};
+ v.constancy = {1};
+ v.divisibility = {highestPow2Divisor(c)};
+ v.knownConstant = c;
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+ setAllPessimistic(op, results);
+ return success();
+ }
+ auto dense = dyn_cast<DenseIntElementsAttr>(op.getValue());
+ if (!dense) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ auto shape = vt.getShape();
if (dense.isSplat()) {
- OffsetsPattern p;
- p.kind = OffsetsPattern::Kind::Broadcast;
- return p;
+ int64_t c = dense.getSplatValue<APInt>().getSExtValue();
+ AxisInfo v = splatAxisInfo(shape, c);
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+
+ // Compute innermost-dim contiguity / constancy / base-divisibility by
+ // iterating the dense values along the inner stride. Outer dims report
+ // pessimistic (1) unless they collapse trivially below.
+ unsigned r = shape.size();
+ int64_t inner = shape.back();
+ int64_t outer = vt.getNumElements() / inner;
+ if (inner < 2 || outer < 1) {
+ // Can't meaningfully analyze a 0/1-element inner dim; fall back to
+ // splat handling already covered, otherwise pessimistic.
+ AxisInfo v = AxisInfo::getPessimistic(r);
+ // For a 1-element inner dim the inner-dim contiguity/constancy is
+ // trivially 1 (already pessimistic).
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
}
- // Check if values form an arithmetic progression.
auto values = llvm::to_vector(dense.getValues<APInt>());
- if (values.size() < 2)
- return std::nullopt;
- int64_t stride =
+ int64_t innerCont = inner;
+ int64_t innerConst = inner;
+ int64_t innerStride =
values[1].getSExtValue() - values[0].getSExtValue();
- for (size_t i = 2; i < values.size(); ++i) {
- int64_t diff = values[i].getSExtValue() - values[i - 1].getSExtValue();
- if (diff != stride)
- return std::nullopt;
- }
- OffsetsPattern p;
- if (stride == 0) {
- p.kind = OffsetsPattern::Kind::Broadcast;
- } else {
- p.kind = OffsetsPattern::Kind::Affine;
- p.stride = stride;
+ int64_t base = values[0].getSExtValue();
+ int64_t baseDiv = highestPow2Divisor(base);
+ for (int64_t o = 0; o < outer; ++o) {
+ int64_t origin = values[o * inner].getSExtValue();
+ baseDiv = std::gcd(baseDiv, highestPow2Divisor(origin));
+ for (int64_t i = 1; i < inner; ++i) {
+ int64_t cur = values[o * inner + i].getSExtValue();
+ int64_t prev = values[o * inner + i - 1].getSExtValue();
+ int64_t diff = cur - prev;
+ if (diff != innerStride)
+ innerStride = std::numeric_limits<int64_t>::min(); // not AP
+ if (diff != 1)
+ innerCont = std::min<int64_t>(innerCont, i);
+ if (diff != 0)
+ innerConst = std::min<int64_t>(innerConst, i);
+ }
}
- return p;
+ AxisInfo v = AxisInfo::getPessimistic(r);
+ if (innerStride == 1)
+ v.contiguity[r - 1] = innerCont;
+ else if (innerStride == 0)
+ v.constancy[r - 1] = innerConst;
+ // For a non-AP inner dim, leave at pessimistic.
+ v.divisibility[r - 1] = baseDiv;
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
}
- // vector.step -> stride 1.
- if (offsets.getDefiningOp<vector::StepOp>()) {
- OffsetsPattern p;
- p.kind = OffsetsPattern::Kind::Affine;
- p.stride = 1;
- return p;
+ // vector.broadcast: source lattice extends to the broadcast dims with
+ // constancy = full extent on those dims. The trailing dims of the source
+ // (if any) align with the trailing dims of the result.
+ LogicalResult visitBroadcast(vector::BroadcastOp op,
+ ArrayRef<const AxisInfoLattice *> operands,
+ ArrayRef<AxisInfoLattice *> results) {
+ auto resTy = dyn_cast<VectorType>(op.getType());
+ if (!resTy) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ unsigned rRank = resTy.getRank();
+ AxisInfo src = operands[0]->getValue();
+ AxisInfo v = AxisInfo::getPessimistic(rRank);
+ auto resShape = resTy.getShape();
+ auto srcVt = dyn_cast<VectorType>(op.getSource().getType());
+ unsigned sRank = srcVt ? srcVt.getRank() : 0;
+ // Broadcast aligns trailing dims of the source with trailing dims of
+ // the result. Leading dims that are 1 in source (or absent) are filled
+ // with constancy = result extent.
+ for (unsigned d = 0; d < rRank; ++d) {
+ int64_t resExt = resShape[d];
+ // Index in source aligned with result dim d, or -1 if d is a
+ // broadcast (front-padded) dim.
+ int sIdx = static_cast<int>(d) - static_cast<int>(rRank - sRank);
+ if (sIdx < 0) {
+ v.constancy[d] = resExt;
+ v.contiguity[d] = 1;
+ v.divisibility[d] =
+ src.isInitialized() ? src.divisibility.front() : 1;
+ continue;
+ }
+ int64_t srcExt = srcVt.getShape()[sIdx];
+ if (srcExt == 1 && resExt > 1) {
+ v.constancy[d] = resExt;
+ v.contiguity[d] = 1;
+ v.divisibility[d] =
+ src.isInitialized() ? src.divisibility[sIdx] : 1;
+ } else if (src.isInitialized()) {
+ v.contiguity[d] = src.contiguity[sIdx];
+ v.constancy[d] = src.constancy[sIdx];
+ v.divisibility[d] = src.divisibility[sIdx];
+ }
+ }
+ if (src.knownConstant)
+ v.knownConstant = src.knownConstant;
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
}
- // Helper to recognize a vector splat with constant integer value.
- auto matchIndexSplat = [](Value v) -> std::optional<int64_t> {
- if (auto bcast = v.getDefiningOp<vector::BroadcastOp>()) {
- APInt c;
- if (matchPattern(bcast.getSource(), m_ConstantInt(&c)))
- return c.getSExtValue();
+ // vector.shape_cast: handle two cases.
+ // (a) Identity-like: shapes match after stripping leading-1 dims —
+ // rebind per-dim info to the new dim positions.
+ // (b) General reshape with the same total element count and row-major
+ // linearization — propagate the source's innermost-dim info (inner
+ // contiguity / constancy) to the destination's innermost dim,
+ // capped by the inner extent. Outer dims stay pessimistic.
+ LogicalResult visitShapeCast(vector::ShapeCastOp op,
+ ArrayRef<const AxisInfoLattice *> operands,
+ ArrayRef<AxisInfoLattice *> results) {
+ auto srcTy = dyn_cast<VectorType>(op.getSource().getType());
+ auto dstTy = dyn_cast<VectorType>(op.getType());
+ if (!srcTy || !dstTy) {
+ setAllPessimistic(op, results);
+ return success();
}
- if (auto cst = v.getDefiningOp<arith::ConstantOp>()) {
- if (auto dense = dyn_cast<DenseIntElementsAttr>(cst.getValue()))
- if (dense.isSplat())
- return dense.getSplatValue<APInt>().getSExtValue();
+ AxisInfo src = operands[0]->getValue();
+ if (!src.isInitialized()) {
+ setAllPessimistic(op, results);
+ return success();
}
- return std::nullopt;
- };
-
- // arith.muli with one splat operand.
- if (auto mul = offsets.getDefiningOp<arith::MulIOp>()) {
- Value lhs = mul.getLhs(), rhs = mul.getRhs();
- auto lhsSplat = matchIndexSplat(lhs);
- auto rhsSplat = matchIndexSplat(rhs);
- Value nonSplat = lhsSplat ? rhs : (rhsSplat ? lhs : Value());
- std::optional<int64_t> factor = lhsSplat ? lhsSplat : rhsSplat;
- if (nonSplat && factor) {
- auto inner = matchOffsetsPattern(nonSplat);
- if (inner && inner->kind == OffsetsPattern::Kind::Affine) {
- OffsetsPattern p;
- p.kind = OffsetsPattern::Kind::Affine;
- p.stride = inner->stride * (*factor);
- return p;
+ // Strip leading 1-dims on both sides; if remaining shapes match, this
+ // is an identity-like reshape.
+ auto stripLeading = [](ArrayRef<int64_t> s) {
+ unsigned i = 0;
+ while (i < s.size() && s[i] == 1)
+ ++i;
+ return s.drop_front(i);
+ };
+ auto sCore = stripLeading(srcTy.getShape());
+ auto dCore = stripLeading(dstTy.getShape());
+ unsigned dRank = dstTy.getRank();
+ AxisInfo v = AxisInfo::getPessimistic(dRank);
+ if (sCore == dCore) {
+ unsigned sLead = srcTy.getRank() - sCore.size();
+ unsigned dLead = dRank - dCore.size();
+ for (unsigned d = dLead; d < dRank; ++d) {
+ unsigned sIdx = sLead + (d - dLead);
+ v.contiguity[d] = src.contiguity[sIdx];
+ v.constancy[d] = src.constancy[sIdx];
+ v.divisibility[d] = src.divisibility[sIdx];
}
- if (inner && inner->kind == OffsetsPattern::Kind::Broadcast) {
- // splat * splat is still uniform.
- return inner;
+ } else {
+ // General linear reshape. Propagate source's inner-dim contiguity /
+ // constancy to dst's inner dim, capped by inner extent. Treat inner
+ // info conservatively as the min across all source dims (so a 1-D
+ // source with full contig => inner-dim contig on dst; an N-D source
+ // collapsed to 1-D inherits the inner-dim info).
+ int64_t innerExt = dstTy.getShape().back();
+ int64_t srcContig = std::numeric_limits<int64_t>::max();
+ int64_t srcConst = std::numeric_limits<int64_t>::max();
+ int64_t srcDiv = src.divisibility[src.getRank() - 1];
+ for (unsigned d = 0; d < src.getRank(); ++d) {
+ srcContig = std::min(srcContig, src.contiguity[d]);
+ srcConst = std::min(srcConst, src.constancy[d]);
}
+ v.contiguity[dRank - 1] = std::min<int64_t>(srcContig, innerExt);
+ v.constancy[dRank - 1] = std::min<int64_t>(srcConst, innerExt);
+ v.divisibility[dRank - 1] = srcDiv;
}
+ if (src.knownConstant)
+ v.knownConstant = src.knownConstant;
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
}
- // arith.addi with a splat operand: stride is unaffected.
- if (auto add = offsets.getDefiningOp<arith::AddIOp>()) {
- Value lhs = add.getLhs(), rhs = add.getRhs();
- auto lhsSplat = matchIndexSplat(lhs);
- auto rhsSplat = matchIndexSplat(rhs);
- Value nonSplat = lhsSplat ? rhs : (rhsSplat ? lhs : Value());
- if (nonSplat)
- return matchOffsetsPattern(nonSplat);
+ template <bool IsSub, typename OpTy>
+ LogicalResult visitAddSub(OpTy op,
+ ArrayRef<const AxisInfoLattice *> operands,
+ ArrayRef<AxisInfoLattice *> results) {
+ auto vt = dyn_cast<VectorType>(op.getType());
+ if (!vt) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ AxisInfo lhs = operands[0]->getValue();
+ AxisInfo rhs = operands[1]->getValue();
+ if (!lhs.isInitialized() || !rhs.isInitialized()) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ unsigned r = vt.getRank();
+ AxisInfo v = AxisInfo::getPessimistic(r);
+ for (unsigned d = 0; d < r; ++d) {
+ int64_t lhsCont = lhs.contiguity[d];
+ int64_t rhsCont = rhs.contiguity[d];
+ int64_t lhsConst = lhs.constancy[d];
+ int64_t rhsConst = rhs.constancy[d];
+ // contiguity propagates through add when one side is constant on the
+ // run, and through sub only when the rhs is constant on the run.
+ int64_t cont = IsSub ? std::min(lhsCont, rhsConst)
+ : std::max(std::min(lhsCont, rhsConst),
+ std::min(rhsCont, lhsConst));
+ v.contiguity[d] = std::max<int64_t>(1, cont);
+ v.constancy[d] = std::min(lhsConst, rhsConst);
+ v.divisibility[d] = std::gcd(lhs.divisibility[d], rhs.divisibility[d]);
+ }
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
}
- return std::nullopt;
+ LogicalResult visitMul(arith::MulIOp op,
+ ArrayRef<const AxisInfoLattice *> operands,
+ ArrayRef<AxisInfoLattice *> results) {
+ auto vt = dyn_cast<VectorType>(op.getType());
+ if (!vt) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ AxisInfo lhs = operands[0]->getValue();
+ AxisInfo rhs = operands[1]->getValue();
+ if (!lhs.isInitialized() || !rhs.isInitialized()) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ unsigned r = vt.getRank();
+ auto shape = vt.getShape();
+ AxisInfo v = AxisInfo::getPessimistic(r);
+ auto unitConstant = [](const AxisInfo &a, unsigned d, int64_t lanes) {
+ return a.knownConstant && *a.knownConstant == 1 &&
+ a.constancy[d] >= lanes;
+ };
+ for (unsigned d = 0; d < r; ++d) {
+ v.constancy[d] = std::min({shape[d], lhs.constancy[d], rhs.constancy[d]});
+ v.divisibility[d] = std::min<int64_t>(
+ kAxisInfoTop, lhs.divisibility[d] * rhs.divisibility[d]);
+ // Multiplying by uniform `s` only keeps contiguity when `s == 1`.
+ if (unitConstant(lhs, d, shape[d]))
+ v.contiguity[d] = std::min(rhs.contiguity[d], shape[d]);
+ else if (unitConstant(rhs, d, shape[d]))
+ v.contiguity[d] = std::min(lhs.contiguity[d], shape[d]);
+ else
+ v.contiguity[d] = 1;
+ }
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+
+ template <typename OpTy>
+ LogicalResult visitPassThrough(OpTy op,
+ ArrayRef<const AxisInfoLattice *> operands,
+ ArrayRef<AxisInfoLattice *> results) {
+ if (!isa<VectorType>(op.getType())) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ propagateIfChanged(results[0], results[0]->join(operands[0]->getValue()));
+ return success();
+ }
+};
+
+} // namespace mlir::xegpu::detail::axis_dataflow
+
+namespace {
+
+using ::mlir::xegpu::detail::axis_dataflow::AxisInfo;
+using ::mlir::xegpu::detail::axis_dataflow::AxisInfoLattice;
+
+//===----------------------------------------------------------------------===//
+// Coalescing decision.
+//===----------------------------------------------------------------------===//
+
+struct CoalesceDecision {
+ enum class Kind { None, Broadcast, Chunked };
+ Kind kind = Kind::None;
+ int64_t factor = 1; // lane_data factor along the innermost dim
+};
+
+/// Largest power-of-two `<= bound` that divides `numLanes`.
+static int64_t largestPow2Divisor(int64_t numLanes, int64_t bound) {
+ if (bound < 2 || numLanes < 2)
+ return 1;
+ int64_t f = std::min<int64_t>(bound, numLanes);
+ // Round down to power of 2.
+ if (!llvm::isPowerOf2_64(f))
+ f = static_cast<int64_t>(llvm::bit_floor(static_cast<uint64_t>(f)));
+ while (f >= 2) {
+ if (numLanes % f == 0)
+ return f;
+ f /= 2;
+ }
+ return 1;
+}
+
+/// Decide how to coalesce given the offsets axis info.
+static CoalesceDecision decide(const AxisInfo &info,
+ ArrayRef<int64_t> offsetsShape,
+ int64_t origChunk, unsigned maxChunkSize) {
+ CoalesceDecision d;
+ if (!info.isInitialized() || offsetsShape.empty())
+ return d;
+ unsigned innerDim = offsetsShape.size() - 1;
+ int64_t inner = offsetsShape[innerDim];
+ if (inner < 2)
+ return d;
+
+ // Broadcast if the innermost dim is uniform across all lanes.
+ if (info.constancy[innerDim] >= inner) {
+ d.kind = CoalesceDecision::Kind::Broadcast;
+ return d;
+ }
+
+ if (origChunk < 1)
+ origChunk = 1;
+ int64_t budget = static_cast<int64_t>(maxChunkSize) / origChunk;
+ if (budget < 2)
+ return d;
+ int64_t bound = std::min<int64_t>(info.contiguity[innerDim], budget);
+ if (bound < 2)
+ return d;
+ int64_t factor = largestPow2Divisor(inner, bound);
+ if (factor < 2)
+ return d;
+ d.kind = CoalesceDecision::Kind::Chunked;
+ d.factor = factor;
+ return d;
}
/// Returns true if `mask` is a constant `dense<true>` vector.
@@ -174,290 +594,153 @@ static bool isAllTrueMask(Value mask) {
return dense.getSplatValue<APInt>().getBoolValue();
}
-/// Compute the largest factor `N` such that:
-/// - `N` divides `numLanes`,
-/// - `N <= maxChunkSize / origChunkSize`,
-/// - lane stride (in elements) equals 1 when scaled by N (i.e. `N * stride
-/// == N * stride`, requires stride to be 1 for true contiguity at the
-/// coalesced granularity).
-///
-/// Returns std::nullopt if no useful coalescing factor exists.
-static std::optional<int64_t> chooseAffineCoalesceFactor(int64_t numLanes,
- int64_t origChunk,
- int64_t stride,
- unsigned maxChunkSize) {
- if (stride != 1)
- return std::nullopt;
- if (numLanes < 2)
- return std::nullopt;
- if (origChunk < 1)
- origChunk = 1;
- int64_t budget = static_cast<int64_t>(maxChunkSize) / origChunk;
- if (budget < 2)
- return std::nullopt;
- // Pick the largest power-of-two factor that divides numLanes and fits
- // budget.
- int64_t factor = 1;
- for (int64_t f = std::min<int64_t>(budget, numLanes); f >= 2; f /= 2) {
- if (numLanes % f == 0) {
- factor = f;
- break;
- }
- }
- if (factor < 2)
- return std::nullopt;
- return factor;
+/// Build a `lane_layout`/`lane_data` layout of rank `rank`, with lane_data
+/// = factor on the innermost dim (1 elsewhere) and lane_layout = inner /
+/// factor on the innermost dim (1 elsewhere).
+static xegpu::LayoutAttr buildLaneDataLayout(MLIRContext *ctx, unsigned rank,
+ int64_t innerLanes,
+ int64_t factor) {
+ SmallVector<int32_t> laneLayout(rank, 1);
+ SmallVector<int32_t> laneData(rank, 1);
+ laneLayout.back() = static_cast<int32_t>(innerLanes / factor);
+ laneData.back() = static_cast<int32_t>(factor);
+ return xegpu::LayoutAttr::get(ctx, laneLayout, laneData);
}
-/// Pick the new offsets value: a sub-vector of length `newLen` extracted from
-/// the original offsets, taking every `factor`-th element. For affine offsets
-/// generated from `vector.step` + scalar adds/muls, the original offsets
-/// already encode the right values; the simplest valid construction is to
-/// rebuild them using the same defining chain at the smaller length.
-///
-/// We take a pragmatic approach: emit
-/// %step = vector.step : vector<newLen x index>
-/// %scaled = arith.muli %step, %splat<factor*stride_element>
-/// %final = arith.addi %scaled, %splat<lane0_offset>
-///
-/// where `lane0_offset` is `offsets[0]` extracted at runtime via
-/// `vector.extract`. For dense-constant offsets we just emit a smaller dense
-/// constant directly.
-static Value buildCoalescedAffineOffsets(OpBuilder &b, Location loc,
- Value origOffsets, int64_t newLen,
- int64_t factor) {
- auto idxTy = b.getIndexType();
- auto newVecTy = VectorType::get({newLen}, idxTy);
-
- // Fast path: dense constant offsets -> emit a smaller dense constant.
- if (auto cst = origOffsets.getDefiningOp<arith::ConstantOp>()) {
- if (auto dense = dyn_cast<DenseIntElementsAttr>(cst.getValue())) {
- auto srcVals = llvm::to_vector(dense.getValues<APInt>());
- SmallVector<APInt> newVals;
- newVals.reserve(newLen);
- for (int64_t i = 0; i < newLen; ++i)
- newVals.push_back(srcVals[i * factor]);
- auto newAttr =
- DenseIntElementsAttr::get(newVecTy, ArrayRef<APInt>(newVals));
- return arith::ConstantOp::create(b, loc, newAttr);
- }
- }
+//===----------------------------------------------------------------------===//
+// Rewrites.
+//===----------------------------------------------------------------------===//
- // General path: rebuild using vector.step + scalar broadcast/muli/addi.
- // Lane-0 offset is the first element of the original offsets.
- Value zero = arith::ConstantIndexOp::create(b, loc, 0);
- Value lane0 =
- vector::ExtractOp::create(b, loc, origOffsets, ArrayRef<int64_t>{0});
- Value step = vector::StepOp::create(b, loc, newVecTy);
- Value factorSplat = arith::ConstantOp::create(
- b, loc,
- DenseIntElementsAttr::get(newVecTy,
- APInt(64, factor, /*isSigned=*/true)));
- // Note: vector.step yields index, factorSplat must also be index-typed.
- // Build factor splat as index using broadcast of scalar.
- factorSplat = vector::BroadcastOp::create(
- b, loc, newVecTy,
- arith::ConstantIndexOp::create(b, loc, factor).getResult());
- Value scaled = arith::MulIOp::create(b, loc, step, factorSplat);
- Value baseSplat = vector::BroadcastOp::create(b, loc, newVecTy, lane0);
- Value result = arith::AddIOp::create(b, loc, scaled, baseSplat);
- (void)zero;
- return result;
+/// Replace an `xegpu.load` whose offsets are uniform along the innermost
+/// dim with a length-1 load + `vector.broadcast` back to the original
+/// value type. Works for any rank; the length-1 load uses an inner-dim
+/// length-1 offsets/mask vector.
+static LogicalResult rewriteBroadcastLoad(xegpu::LoadGatherOp op,
+ PatternRewriter &rewriter) {
+ Location loc = op.getLoc();
+ auto valueTy = op.getValueType();
+ auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
+ if (!valueTy || !offsetsTy)
+ return failure();
+
+ // Extract a scalar offset from index 0...0.
+ SmallVector<int64_t> zeros(offsetsTy.getRank(), 0);
+ Value scalarOffset =
+ vector::ExtractOp::create(rewriter, loc, op.getOffsets(), zeros);
+ auto idxVecTy = VectorType::get({1}, rewriter.getIndexType());
+ auto maskVecTy = VectorType::get({1}, rewriter.getI1Type());
+ Value newOffsets =
+ vector::BroadcastOp::create(rewriter, loc, idxVecTy, scalarOffset);
+ Value newMask = arith::ConstantOp::create(
+ rewriter, loc, DenseIntElementsAttr::get(maskVecTy, true));
+ auto newValueTy = VectorType::get({1}, valueTy.getElementType());
+ auto newLoad = xegpu::LoadGatherOp::create(
+ rewriter, loc, newValueTy, op.getSource(), newOffsets, newMask,
+ /*chunk_size=*/IntegerAttr(), op.getL1HintAttr(), op.getL2HintAttr(),
+ op.getL3HintAttr(), /*layout=*/xegpu::DistributeLayoutAttr());
+ Value scalar = vector::ExtractOp::create(rewriter, loc, newLoad.getResult(),
+ ArrayRef<int64_t>{0});
+ Value bcast = vector::BroadcastOp::create(rewriter, loc, valueTy, scalar);
+ rewriter.replaceOp(op, bcast);
+ return success();
}
-/// Build a `dense<true>` mask vector of length `newLen`.
-static Value buildAllTrueMask(OpBuilder &b, Location loc, int64_t newLen) {
- auto vecTy = VectorType::get({newLen}, b.getI1Type());
- auto attr = DenseIntElementsAttr::get(vecTy, true);
- return arith::ConstantOp::create(b, loc, attr);
-}
+namespace {
-/// Pattern that coalesces a `xegpu.load` op.
struct CoalesceLoadPattern final : OpRewritePattern<xegpu::LoadGatherOp> {
- CoalesceLoadPattern(MLIRContext *ctx, unsigned maxChunkSize)
- : OpRewritePattern(ctx), maxChunkSize(maxChunkSize) {}
+ CoalesceLoadPattern(MLIRContext *ctx, unsigned maxChunkSize,
+ DataFlowSolver &solver)
+ : OpRewritePattern(ctx), maxChunkSize(maxChunkSize), solver(solver) {}
LogicalResult matchAndRewrite(xegpu::LoadGatherOp op,
PatternRewriter &rewriter) const override {
- // Only the source-as-pointer/memref form has a vector offsets operand we
- // can analyze; the tensor_desc form has no offsets here.
auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
- if (!offsetsTy || offsetsTy.getRank() != 1)
- return rewriter.notifyMatchFailure(op, "expected 1-D vector offsets");
+ if (!offsetsTy)
+ return rewriter.notifyMatchFailure(op, "expected vector offsets");
if (offsetsTy.getNumElements() <= 1)
return rewriter.notifyMatchFailure(op, "nothing to coalesce");
auto valueTy = op.getValueType();
- if (!valueTy || valueTy.getRank() != 1)
- return rewriter.notifyMatchFailure(op, "expected 1-D vector value");
+ if (!valueTy)
+ return rewriter.notifyMatchFailure(op, "expected vector value");
if (!isAllTrueMask(op.getMask()))
return rewriter.notifyMatchFailure(op, "non-uniform mask");
- auto pattern = matchOffsetsPattern(op.getOffsets());
- if (!pattern)
- return rewriter.notifyMatchFailure(op, "offsets not coalescible");
+ // Already coalesced (a previous run, or another pass tagged it).
+ if (auto layout = op.getLayoutAttr())
+ if (!layout.getEffectiveLaneDataAsInt().empty())
+ return rewriter.notifyMatchFailure(op, "lane_data already set");
- int64_t numLanes = offsetsTy.getNumElements();
- int64_t origChunk =
- static_cast<int64_t>(op.getChunkSize().value_or(1));
-
- Location loc = op.getLoc();
-
- if (pattern->kind == OffsetsPattern::Kind::Broadcast) {
- // All lanes load the same address: emit a single scalar load and
- // broadcast.
- Value scalarOffset = vector::ExtractOp::create(
- rewriter, loc, op.getOffsets(), ArrayRef<int64_t>{0});
- Value scalarMask = arith::ConstantOp::create(
- rewriter, loc, rewriter.getOneAttr(rewriter.getI1Type()));
- // Result of the scalar load matches the original element type with the
- // chunk dimension preserved if any (1-D value, length = chunk_size or 1).
- int64_t scalarLen = numLanes; // value length per lane is implicit; for
- // 1-D the whole vector represents lanes.
- // Build a length-1 offsets vector to keep types consistent (the op
- // accepts scalar offsets too, but our incoming form is vector).
- auto idxVecTy = VectorType::get({1}, rewriter.getIndexType());
- auto maskVecTy = VectorType::get({1}, rewriter.getI1Type());
- Value newOffsets = vector::BroadcastOp::create(
- rewriter, loc, idxVecTy, scalarOffset);
- Value newMask = arith::ConstantOp::create(
- rewriter, loc, DenseIntElementsAttr::get(maskVecTy, true));
- auto newValueTy =
- VectorType::get({1}, valueTy.getElementType());
- auto newLoad = xegpu::LoadGatherOp::create(
- rewriter, loc, newValueTy, op.getSource(), newOffsets, newMask,
- /*chunk_size=*/IntegerAttr(), op.getL1HintAttr(),
- op.getL2HintAttr(), op.getL3HintAttr(),
- /*layout=*/xegpu::DistributeLayoutAttr());
- // Splat to original shape.
- Value scalar = vector::ExtractOp::create(
- rewriter, loc, newLoad.getResult(), ArrayRef<int64_t>{0});
- Value bcast =
- vector::BroadcastOp::create(rewriter, loc, valueTy, scalar);
- rewriter.replaceOp(op, bcast);
- (void)scalarMask;
- (void)scalarLen;
- return success();
- }
+ const auto *lat = solver.lookupState<AxisInfoLattice>(op.getOffsets());
+ if (!lat || !lat->getValue().isInitialized())
+ return rewriter.notifyMatchFailure(op, "no axis-info available");
- // Affine path.
- auto factorOpt = chooseAffineCoalesceFactor(numLanes, origChunk,
- pattern->stride, maxChunkSize);
- if (!factorOpt)
- return rewriter.notifyMatchFailure(op, "no useful coalesce factor");
- int64_t factor = *factorOpt;
- int64_t newLanes = numLanes / factor;
- int64_t newChunk = origChunk * factor;
-
- Value newOffsets = buildCoalescedAffineOffsets(rewriter, loc,
- op.getOffsets(),
- newLanes, factor);
- Value newMask = buildAllTrueMask(rewriter, loc, newLanes);
-
- // The verifier requires `chunk_size > 1` to be paired with a 2-D value
- // vector of shape `<lanes x chunk>`. Build that 2-D type and shape_cast
- // back to the original 1-D shape for the consumer.
- SmallVector<int64_t> newShape;
- if (valueTy.getRank() == 1) {
- newShape = {newLanes, newChunk};
- } else {
- newShape = llvm::to_vector(valueTy.getShape());
- newShape[0] = newLanes;
- newShape.back() = newChunk;
- }
- auto newValueTy = VectorType::get(newShape, valueTy.getElementType());
+ int64_t origChunk = static_cast<int64_t>(op.getChunkSize().value_or(1));
+ auto d = decide(lat->getValue(), offsetsTy.getShape(), origChunk,
+ maxChunkSize);
- auto newChunkAttr = rewriter.getI64IntegerAttr(newChunk);
- auto newLoad = xegpu::LoadGatherOp::create(
- rewriter, loc, newValueTy, op.getSource(), newOffsets, newMask,
- newChunkAttr, op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), /*layout=*/xegpu::DistributeLayoutAttr());
+ if (d.kind == CoalesceDecision::Kind::Broadcast)
+ return rewriteBroadcastLoad(op, rewriter);
- if (newValueTy == valueTy) {
- rewriter.replaceOp(op, newLoad.getResult());
- } else {
- Value reshaped = vector::ShapeCastOp::create(
- rewriter, loc, valueTy, newLoad.getResult());
- rewriter.replaceOp(op, reshaped);
- }
+ if (d.kind != CoalesceDecision::Kind::Chunked)
+ return rewriter.notifyMatchFailure(op, "offsets not coalescible");
+
+ int64_t innerLanes = offsetsTy.getShape().back();
+ auto layout = buildLaneDataLayout(op.getContext(), valueTy.getRank(),
+ innerLanes, d.factor);
+ rewriter.modifyOpInPlace(op, [&] { op.setLayoutAttr(layout); });
return success();
}
unsigned maxChunkSize;
+ DataFlowSolver &solver;
};
-/// Pattern that coalesces a `xegpu.store` op.
struct CoalesceStorePattern final : OpRewritePattern<xegpu::StoreScatterOp> {
- CoalesceStorePattern(MLIRContext *ctx, unsigned maxChunkSize)
- : OpRewritePattern(ctx), maxChunkSize(maxChunkSize) {}
+ CoalesceStorePattern(MLIRContext *ctx, unsigned maxChunkSize,
+ DataFlowSolver &solver)
+ : OpRewritePattern(ctx), maxChunkSize(maxChunkSize), solver(solver) {}
LogicalResult matchAndRewrite(xegpu::StoreScatterOp op,
PatternRewriter &rewriter) const override {
auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
- if (!offsetsTy || offsetsTy.getRank() != 1)
- return rewriter.notifyMatchFailure(op, "expected 1-D vector offsets");
+ if (!offsetsTy)
+ return rewriter.notifyMatchFailure(op, "expected vector offsets");
if (offsetsTy.getNumElements() <= 1)
return rewriter.notifyMatchFailure(op, "nothing to coalesce");
auto valueTy = op.getValueType();
- if (!valueTy || valueTy.getRank() != 1)
- return rewriter.notifyMatchFailure(op, "expected 1-D vector value");
+ if (!valueTy)
+ return rewriter.notifyMatchFailure(op, "expected vector value");
if (!isAllTrueMask(op.getMask()))
return rewriter.notifyMatchFailure(op, "non-uniform mask");
- auto pattern = matchOffsetsPattern(op.getOffsets());
- if (!pattern)
- return rewriter.notifyMatchFailure(op, "offsets not coalescible");
+ if (auto layout = op.getLayoutAttr())
+ if (!layout.getEffectiveLaneDataAsInt().empty())
+ return rewriter.notifyMatchFailure(op, "lane_data already set");
+
+ const auto *lat = solver.lookupState<AxisInfoLattice>(op.getOffsets());
+ if (!lat || !lat->getValue().isInitialized())
+ return rewriter.notifyMatchFailure(op, "no axis-info available");
+
+ int64_t origChunk = static_cast<int64_t>(op.getChunkSize().value_or(1));
+ auto d = decide(lat->getValue(), offsetsTy.getShape(), origChunk,
+ maxChunkSize);
- // For the broadcast case, all lanes write the same address. That is not a
- // meaningful coalesce target (last writer wins, semantics-preserving
- // rewrite would need a single arbitrary lane to win). Skip.
- if (pattern->kind == OffsetsPattern::Kind::Broadcast)
+ if (d.kind == CoalesceDecision::Kind::Broadcast)
return rewriter.notifyMatchFailure(
op, "all-equal offsets on store would be ambiguous");
- int64_t numLanes = offsetsTy.getNumElements();
- int64_t origChunk =
- static_cast<int64_t>(op.getChunkSize().value_or(1));
- auto factorOpt = chooseAffineCoalesceFactor(numLanes, origChunk,
- pattern->stride, maxChunkSize);
- if (!factorOpt)
- return rewriter.notifyMatchFailure(op, "no useful coalesce factor");
- int64_t factor = *factorOpt;
- int64_t newLanes = numLanes / factor;
- int64_t newChunk = origChunk * factor;
-
- Location loc = op.getLoc();
- Value newOffsets = buildCoalescedAffineOffsets(rewriter, loc,
- op.getOffsets(),
- newLanes, factor);
- Value newMask = buildAllTrueMask(rewriter, loc, newLanes);
-
- auto newChunkAttr = rewriter.getI64IntegerAttr(newChunk);
- // The verifier requires `chunk_size > 1` to be paired with a 2-D value
- // vector of shape `<lanes x chunk>`; shape_cast the original 1-D value
- // into that shape.
- SmallVector<int64_t> newShape;
- if (valueTy.getRank() == 1)
- newShape = {newLanes, newChunk};
- else {
- newShape = llvm::to_vector(valueTy.getShape());
- newShape[0] = newLanes;
- newShape.back() = newChunk;
- }
- auto newValTy = VectorType::get(newShape, valueTy.getElementType());
- Value newValue = op.getValue();
- if (newValTy != valueTy)
- newValue =
- vector::ShapeCastOp::create(rewriter, loc, newValTy, op.getValue());
- xegpu::StoreScatterOp::create(rewriter, loc, newValue, op.getDest(),
- newOffsets, newMask, newChunkAttr,
- op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(),
- /*layout=*/xegpu::DistributeLayoutAttr());
- rewriter.eraseOp(op);
+ if (d.kind != CoalesceDecision::Kind::Chunked)
+ return rewriter.notifyMatchFailure(op, "offsets not coalescible");
+
+ int64_t innerLanes = offsetsTy.getShape().back();
+ auto layout = buildLaneDataLayout(op.getContext(), valueTy.getRank(),
+ innerLanes, d.factor);
+ rewriter.modifyOpInPlace(op, [&] { op.setLayoutAttr(layout); });
return success();
}
unsigned maxChunkSize;
+ DataFlowSolver &solver;
};
struct XeGPUCoalesceGatherScatterPass final
@@ -466,12 +749,23 @@ struct XeGPUCoalesceGatherScatterPass final
using XeGPUCoalesceGatherScatterBase::XeGPUCoalesceGatherScatterBase;
void runOnOperation() override {
+ Operation *root = getOperation();
+
+ DataFlowSolver solver;
+ solver.load<dataflow::DeadCodeAnalysis>();
+ solver.load<mlir::xegpu::detail::axis_dataflow::AxisInfoAnalysis>();
+ if (failed(solver.initializeAndRun(root)))
+ return signalPassFailure();
+
MLIRContext *ctx = &getContext();
RewritePatternSet patterns(ctx);
- patterns.add<CoalesceLoadPattern, CoalesceStorePattern>(ctx, maxChunkSize);
- if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+ patterns.add<CoalesceLoadPattern, CoalesceStorePattern>(ctx, maxChunkSize,
+ solver);
+ if (failed(applyPatternsGreedily(root, std::move(patterns))))
return signalPassFailure();
}
};
} // namespace
+
+} // namespace
diff --git a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
index be205f152604e..6f34618f386e4 100644
--- a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
+++ b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
@@ -3,14 +3,13 @@
// -----
// vector.step offsets -> stride 1, fully coalescible.
-// 32 lanes, max-chunk-size = 8 -> factor 8 -> 4 lanes, chunk_size = 8.
-// CHECK-LABEL: @load_step_offsets
-// CHECK-SAME: (%[[ARG:.*]]: i64) -> vector<32xf32>
-// CHECK: %[[STEP:.*]] = vector.step : vector<4xindex>
-// CHECK: %[[LOAD:.*]] = xegpu.load %[[ARG]][{{.*}}], {{.*}} <{chunk_size = 8 : i64}>
-// CHECK-SAME: : i64, vector<4xindex>, vector<4xi1> -> vector<4x8xf32>
-// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] : vector<4x8xf32> to vector<32xf32>
-// CHECK: return %[[CAST]]
+// 32 lanes, max-chunk-size = 8 -> factor 8 -> lane_layout=[4], lane_data=[8].
+// CHECK-LABEL: func.func @load_step_offsets(
+// CHECK: %[[STEP:.*]] = vector.step : vector<32xindex>
+// CHECK: %[[LOAD:.*]] = xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<lane_layout = [4], lane_data = [8]>
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+// CHECK: return %[[LOAD]]
func.func @load_step_offsets(%ptr: i64) -> vector<32xf32> {
%offsets = vector.step : vector<32xindex>
%mask = arith.constant dense<true> : vector<32xi1>
@@ -20,11 +19,11 @@ func.func @load_step_offsets(%ptr: i64) -> vector<32xf32> {
}
// -----
-// max-chunk-size = 4: factor 4 -> 8 lanes, chunk_size = 4.
-// CHECK4-LABEL: @load_step_offsets_chunk4
-// CHECK4: xegpu.load {{.*}} <{chunk_size = 4 : i64}>
-// CHECK4-SAME: : i64, vector<8xindex>, vector<8xi1> -> vector<8x4xf32>
-// CHECK4: vector.shape_cast {{.*}} : vector<8x4xf32> to vector<32xf32>
+// max-chunk-size = 4: factor 4 -> lane_layout=[8], lane_data=[4].
+// CHECK4-LABEL: func.func @load_step_offsets_chunk4(
+// CHECK4: xegpu.load
+// CHECK4-SAME: layout = #xegpu.layout<lane_layout = [8], lane_data = [4]>
+// CHECK4-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
func.func @load_step_offsets_chunk4(%ptr: i64) -> vector<32xf32> {
%offsets = vector.step : vector<32xindex>
%mask = arith.constant dense<true> : vector<32xi1>
@@ -35,11 +34,10 @@ func.func @load_step_offsets_chunk4(%ptr: i64) -> vector<32xf32> {
// -----
// Dense constant arithmetic progression with stride 1.
-// CHECK-LABEL: @load_dense_ap_offsets
-// CHECK: %[[CST:.*]] = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
-// CHECK: xegpu.load {{.*}}[%[[CST]]], {{.*}} <{chunk_size = 8 : i64}>
-// CHECK-SAME: : i64, vector<4xindex>, vector<4xi1> -> vector<4x8xi32>
-// CHECK: vector.shape_cast {{.*}} : vector<4x8xi32> to vector<32xi32>
+// CHECK-LABEL: func.func @load_dense_ap_offsets(
+// CHECK: xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<lane_layout = [4], lane_data = [8]>
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xi32>
func.func @load_dense_ap_offsets(%ptr: i64) -> vector<32xi32> {
%offsets = arith.constant dense<[
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
@@ -53,7 +51,7 @@ func.func @load_dense_ap_offsets(%ptr: i64) -> vector<32xi32> {
// -----
// All-equal offsets -> broadcast load (single load + vector.broadcast).
-// CHECK-LABEL: @load_broadcast_offsets
+// CHECK-LABEL: func.func @load_broadcast_offsets(
// CHECK: %[[L:.*]] = xegpu.load
// CHECK-SAME: : i64, vector<1xindex>, vector<1xi1> -> vector<1xf32>
// CHECK: %[[E:.*]] = vector.extract %[[L]][0]
@@ -68,9 +66,10 @@ func.func @load_broadcast_offsets(%ptr: i64) -> vector<32xf32> {
}
// -----
-// Non-stride-1 (stride 4) offsets: not contiguous, not coalesced.
-// CHECK-LABEL: @load_stride4_unchanged
+// Non-stride-1 (stride 4) offsets: not contiguous, no layout attached.
+// CHECK-LABEL: func.func @load_stride4_unchanged(
// CHECK: xegpu.load
+// CHECK-NOT: lane_data
// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
func.func @load_stride4_unchanged(%ptr: i64) -> vector<32xf32> {
%c4 = arith.constant 4 : index
@@ -85,8 +84,9 @@ func.func @load_stride4_unchanged(%ptr: i64) -> vector<32xf32> {
// -----
// Non-uniform mask: not coalesced.
-// CHECK-LABEL: @load_partial_mask_unchanged
+// CHECK-LABEL: func.func @load_partial_mask_unchanged(
// CHECK: xegpu.load
+// CHECK-NOT: lane_data
// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
func.func @load_partial_mask_unchanged(%ptr: i64, %mask: vector<32xi1>) -> vector<32xf32> {
%offsets = vector.step : vector<32xindex>
@@ -97,11 +97,10 @@ func.func @load_partial_mask_unchanged(%ptr: i64, %mask: vector<32xi1>) -> vecto
// -----
// Store with vector.step offsets coalesces.
-// CHECK-LABEL: @store_step_offsets
-// CHECK: vector.shape_cast {{.*}} : vector<32xf32> to vector<4x8xf32>
+// CHECK-LABEL: func.func @store_step_offsets(
// CHECK: xegpu.store
-// CHECK-SAME: <{chunk_size = 8 : i64}>
-// CHECK-SAME: vector<4x8xf32>, i64, vector<4xindex>, vector<4xi1>
+// CHECK-SAME: layout = #xegpu.layout<lane_layout = [4], lane_data = [8]>
+// CHECK-SAME: : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
func.func @store_step_offsets(%ptr: i64, %v: vector<32xf32>) {
%offsets = vector.step : vector<32xindex>
%mask = arith.constant dense<true> : vector<32xi1>
@@ -112,8 +111,9 @@ func.func @store_step_offsets(%ptr: i64, %v: vector<32xf32>) {
// -----
// Store with all-equal offsets is left alone (ambiguous semantics).
-// CHECK-LABEL: @store_broadcast_offsets_unchanged
+// CHECK-LABEL: func.func @store_broadcast_offsets_unchanged(
// CHECK: xegpu.store
+// CHECK-NOT: lane_data
// CHECK-SAME: vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
func.func @store_broadcast_offsets_unchanged(%ptr: i64, %v: vector<32xf32>) {
%offsets = arith.constant dense<0> : vector<32xindex>
@@ -125,11 +125,10 @@ func.func @store_broadcast_offsets_unchanged(%ptr: i64, %v: vector<32xf32>) {
// -----
// memref-source variant of load coalesces too.
-// CHECK-LABEL: @load_memref_step
+// CHECK-LABEL: func.func @load_memref_step(
// CHECK: xegpu.load
-// CHECK-SAME: <{chunk_size = 8 : i64}>
-// CHECK-SAME: : memref<1024xf32>, vector<4xindex>, vector<4xi1> -> vector<4x8xf32>
-// CHECK: vector.shape_cast {{.*}} : vector<4x8xf32> to vector<32xf32>
+// CHECK-SAME: layout = #xegpu.layout<lane_layout = [4], lane_data = [8]>
+// CHECK-SAME: : memref<1024xf32>, vector<32xindex>, vector<32xi1> -> vector<32xf32>
func.func @load_memref_step(%m: memref<1024xf32>) -> vector<32xf32> {
%offsets = vector.step : vector<32xindex>
%mask = arith.constant dense<true> : vector<32xi1>
@@ -137,3 +136,87 @@ func.func @load_memref_step(%m: memref<1024xf32>) -> vector<32xf32> {
: memref<1024xf32>, vector<32xindex>, vector<32xi1> -> vector<32xf32>
return %v : vector<32xf32>
}
+
+// -----
+// 2-D offsets with leading unit dim: inner dim treated as lane dim.
+// vector<1x32xindex> stride-1 -> lane_data=[1, 8], lane_layout=[1, 4].
+// CHECK-LABEL: func.func @load_2d_leading_unit(
+// CHECK: xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<lane_layout = [1, 4], lane_data = [1, 8]>
+// CHECK-SAME: : i64, vector<1x32xindex>, vector<1x32xi1> -> vector<1x32xf32>
+func.func @load_2d_leading_unit(%ptr: i64) -> vector<1x32xf32> {
+ %step = vector.step : vector<32xindex>
+ %offsets = vector.shape_cast %step : vector<32xindex> to vector<1x32xindex>
+ %mask = arith.constant dense<true> : vector<1x32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<1x32xindex>, vector<1x32xi1> -> vector<1x32xf32>
+ return %v : vector<1x32xf32>
+}
+
+// -----
+// True 2-D dense AP: each row stride-1, 16 lanes per row -> lane_data=[1,8],
+// lane_layout=[1, 2].
+// CHECK-LABEL: func.func @load_2d_dense_ap(
+// CHECK: xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<lane_layout = [1, 2], lane_data = [1, 8]>
+// CHECK-SAME: : i64, vector<2x16xindex>, vector<2x16xi1> -> vector<2x16xf32>
+func.func @load_2d_dense_ap(%ptr: i64) -> vector<2x16xf32> {
+ %offsets = arith.constant dense<[
+ [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
+ [16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]
+ ]> : vector<2x16xindex>
+ %mask = arith.constant dense<true> : vector<2x16xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<2x16xindex>, vector<2x16xi1> -> vector<2x16xf32>
+ return %v : vector<2x16xf32>
+}
+
+// -----
+// 1-D step reshape_cast'ed to 4x8: inner extent 8, factor 8
+// -> lane_layout=[1, 1], lane_data=[1, 8] (each of 8 inner positions is a
+// lane that fetches a contiguous chunk of 8).
+// CHECK-LABEL: func.func @load_4x8_step_shape_cast(
+// CHECK: xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<lane_layout = [1, 1], lane_data = [1, 8]>
+// CHECK-SAME: : i64, vector<4x8xindex>, vector<4x8xi1> -> vector<4x8xf32>
+func.func @load_4x8_step_shape_cast(%ptr: i64) -> vector<4x8xf32> {
+ %step = vector.step : vector<32xindex>
+ %offsets = vector.shape_cast %step : vector<32xindex> to vector<4x8xindex>
+ %mask = arith.constant dense<true> : vector<4x8xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<4x8xindex>, vector<4x8xi1> -> vector<4x8xf32>
+ return %v : vector<4x8xf32>
+}
+
+// -----
+// Negative: 2-D dense values where inner row is not stride-1 AP. Should
+// not coalesce.
+// CHECK-LABEL: func.func @load_2d_non_ap_unchanged(
+// CHECK: xegpu.load
+// CHECK-NOT: lane_data
+// CHECK-SAME: : i64, vector<2x16xindex>, vector<2x16xi1> -> vector<2x16xf32>
+func.func @load_2d_non_ap_unchanged(%ptr: i64) -> vector<2x16xf32> {
+ %offsets = arith.constant dense<[
+ [0, 1, 2, 3, 4, 5, 6, 7, 100, 9, 10, 11, 12, 13, 14, 15],
+ [16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]
+ ]> : vector<2x16xindex>
+ %mask = arith.constant dense<true> : vector<2x16xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<2x16xindex>, vector<2x16xi1> -> vector<2x16xf32>
+ return %v : vector<2x16xf32>
+}
+
+// -----
+// 2-D store with step + shape_cast offsets coalesces too.
+// CHECK-LABEL: func.func @store_2d_step(
+// CHECK: xegpu.store
+// CHECK-SAME: layout = #xegpu.layout<lane_layout = [1, 4], lane_data = [1, 8]>
+// CHECK-SAME: : vector<1x32xf32>, i64, vector<1x32xindex>, vector<1x32xi1>
+func.func @store_2d_step(%ptr: i64, %v: vector<1x32xf32>) {
+ %step = vector.step : vector<32xindex>
+ %offsets = vector.shape_cast %step : vector<32xindex> to vector<1x32xindex>
+ %mask = arith.constant dense<true> : vector<1x32xi1>
+ xegpu.store %v, %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : vector<1x32xf32>, i64, vector<1x32xindex>, vector<1x32xi1>
+ return
+}
>From cff20c8427b524122f26bdd6c2a2f9690cf26c28 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Thu, 28 May 2026 21:06:12 +0000
Subject: [PATCH 3/5] [mlir][XeGPU][GPU] Wire coalesce pass into
GPUToXeVMPipeline; pick lane_layout via subgroup size; extend AxisInfo with
innerStride and divui/remui/divsi/remsi/andi/shli/shrui/select/transpose
visitors.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Multiple related changes to `xegpu-coalesce-gather-scatter` and the GPU
pipeline:
1. **GPUToXeVMPipeline**: schedule `xegpu-coalesce-gather-scatter` after
`xegpu-wg-to-sg-distribute` and before `xegpu-propagate-layout("inst")`
so the coalesced shape participates in `inst_data` selection.
2. **Lane-layout-first decision**. The coalesce decision picks
`lane_layout[inner] = min(subgroupSize, inner)` first (the same default
rule as `xegpu-propagate-layout`), then derives `lane_data[inner]` from
the per-lane budget × contiguity × `maxChunkSize / origChunkSize`. The
subgroup size is read per-op from the enclosing `gpu.module`'s
`xevm.target` chip via `xegpu::getChipStr` / `xegpu::uArch::getUArch`,
falling back to 16 when no chip is found (so plain unit lit tests still
work).
3. **inst_data on emitted layouts**. New layouts include
`inst_data = lane_layout * lane_data` so the per-dim invariant
`inst_data[d] == lane_layout[d] * lane_data[d]` holds for downstream
blocking.
4. **chunk_size handling**.
- Skip the rewrite when the op already declares an explicit
`chunk_size > 1`: the verifier-imposed value/mask shape relationship
makes coalescing unsafe without a deeper rewrite.
- On a successful coalesce, drop a trivial `chunk_size = 1` attribute
when the new `lane_data` FCD > 1 (the new layout subsumes it).
5. **AxisInfo lattice extension**. Add an optional `innerStride` slot to
the per-value lattice (the AP stride along the innermost dim). It is
populated by `visitStep` (= 1), splat constants (= 0), and dense AP
constants. `visitBroadcast`, `visitShapeCast`, `visitTranspose`,
`visitAddSub`, and `visitMul` propagate it where the algebra survives.
6. **New transfer functions**:
- `vector.transpose` — permutes per-dim contiguity / constancy /
divisibility according to the transpose permutation; preserves
`innerStride` only when the new inner dim came from the old inner
dim. This is what lets the classic `transpose(broadcast(rows)) +
broadcast(step)` 2-D scatter pattern keep inner-dim contiguity
through the trailing `addi`.
- `arith.divui` / `arith.divsi` / `arith.remui` / `arith.remsi` by a
uniform positive constant `c`. Division: when `c | innerStride` and
`c | divisibility[inner]`, set new stride to `s / c` and update
inner-dim divisibility. Remainder: when `c | innerStride`, set
`innerStride = 0` (inner-dim constant).
- `arith.andi` with a uniform constant mask: handles `m == 0` (zero)
and `m == P-1` for power-of-two `P` (equivalent to `mod P`).
- `arith.shli` / `arith.shrui` by a uniform constant `k`: reduces to
mul / div by `1 << k`.
- `arith.select`: result is `AxisInfo::join` of the two arms.
Lit tests are reworked to match the lane-layout-first policy and the new
`inst_data` field, and add coverage for each new transfer function plus
the `chunk_size > 1` skip and the `chunk_size = 1` drop-on-success
behavior.
Co-Authored-By: Claude Opus 4.7 <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/Transforms/Passes.td | 9 +
.../GPU/Pipelines/GPUToXeVMPipeline.cpp | 2 +
.../Transforms/XeGPUCoalesceGatherScatter.cpp | 454 +++++++++++++++++-
.../XeGPU/coalesce-gather-scatter.mlir | 324 ++++++++++++-
4 files changed, 740 insertions(+), 49 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
index 1ac15cc80f4e7..caf33cc1e13c9 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
@@ -145,6 +145,15 @@ def XeGPUCoalesceGatherScatter : Pass<"xegpu-coalesce-gather-scatter"> {
Pass options:
- `max-chunk-size`: upper bound on the produced `chunk_size`. Defaults to
8, matching typical Xe scatter chunk-size limits.
+
+ Ops that already declare an explicit `chunk_size > 1` are left alone:
+ a downstream pass has already committed to a per-lane chunked access,
+ and the verifier-imposed shape relationship between value, offsets,
+ and mask makes coalescing unsafe without a deeper rewrite.
+
+ The subgroup size is read per-op from the enclosing `gpu.module`'s
+ `xevm.target` chip (the same lookup as `xegpu-propagate-layout`). If no
+ target chip is found, the pass falls back to a subgroup size of 16.
}];
let dependentDialects = [
"arith::ArithDialect", "memref::MemRefDialect", "xegpu::XeGPUDialect",
diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
index 7600ec39fb3f5..c515721363a73 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
@@ -74,6 +74,8 @@ void buildGPUPassPipeline(OpPassManager &pm,
pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
pm.addNestedPass<gpu::GPUModuleOp>(createLowerAffinePass());
pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
+ pm.addNestedPass<gpu::GPUModuleOp>(
+ xegpu::createXeGPUCoalesceGatherScatter());
xegpu::XeGPUPropagateLayoutOptions instDataOptions;
instDataOptions.layoutKind = "inst";
pm.addNestedPass<gpu::GPUModuleOp>(
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
index 91808b711cf12..4320ccb2aa6bb 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
@@ -37,6 +37,8 @@
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
+#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
+#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Matchers.h"
@@ -83,13 +85,20 @@ static constexpr int64_t kAxisInfoTop = 1LL << 30;
/// dimension `d`.
/// - `knownConstant`: scalar value if the entire vector is uniformly
/// known to be a single constant.
+/// - `innerStride`: when set, every "row" along the innermost dimension
+/// is an arithmetic progression with this stride. Per-row base may
+/// differ across outer indices; per-row alignment is captured by
+/// `divisibility[innerDim]`. `innerStride = 1` implies stride-1
+/// contiguity; `innerStride = 0` implies inner-dim constancy.
///
-/// Pessimistic / entry value: contiguity=1, constancy=1, divisibility=1.
+/// Pessimistic / entry value: contiguity=1, constancy=1, divisibility=1,
+/// innerStride absent.
struct AxisInfo {
SmallVector<int64_t> contiguity;
SmallVector<int64_t> constancy;
SmallVector<int64_t> divisibility;
std::optional<int64_t> knownConstant;
+ std::optional<int64_t> innerStride;
AxisInfo() = default;
@@ -107,7 +116,7 @@ struct AxisInfo {
bool operator==(const AxisInfo &rhs) const {
return contiguity == rhs.contiguity && constancy == rhs.constancy &&
divisibility == rhs.divisibility &&
- knownConstant == rhs.knownConstant;
+ knownConstant == rhs.knownConstant && innerStride == rhs.innerStride;
}
/// Conservative join. Two values reaching the same SSA value via different
@@ -131,6 +140,9 @@ struct AxisInfo {
if (lhs.knownConstant && rhs.knownConstant &&
*lhs.knownConstant == *rhs.knownConstant)
out.knownConstant = lhs.knownConstant;
+ if (lhs.innerStride && rhs.innerStride &&
+ *lhs.innerStride == *rhs.innerStride)
+ out.innerStride = lhs.innerStride;
return out;
}
@@ -144,6 +156,8 @@ struct AxisInfo {
os << "]";
if (knownConstant)
os << " const=" << *knownConstant;
+ if (innerStride)
+ os << " innerStride=" << *innerStride;
}
};
@@ -172,6 +186,7 @@ static AxisInfo splatAxisInfo(ArrayRef<int64_t> shape, int64_t c) {
v.constancy.assign(shape.begin(), shape.end());
v.divisibility.assign(r, highestPow2Divisor(c));
v.knownConstant = c;
+ v.innerStride = 0;
return v;
}
@@ -194,12 +209,34 @@ class AxisInfoAnalysis
return visitBroadcast(bcast, operands, results);
if (auto sc = dyn_cast<vector::ShapeCastOp>(op))
return visitShapeCast(sc, operands, results);
+ if (auto tp = dyn_cast<vector::TransposeOp>(op))
+ return visitTranspose(tp, operands, results);
if (auto add = dyn_cast<arith::AddIOp>(op))
return visitAddSub</*IsSub=*/false>(add, operands, results);
if (auto sub = dyn_cast<arith::SubIOp>(op))
return visitAddSub</*IsSub=*/true>(sub, operands, results);
if (auto mul = dyn_cast<arith::MulIOp>(op))
return visitMul(mul, operands, results);
+ if (auto div = dyn_cast<arith::DivUIOp>(op))
+ return visitDivRem</*IsSigned=*/false, /*IsRem=*/false>(div, operands,
+ results);
+ if (auto div = dyn_cast<arith::DivSIOp>(op))
+ return visitDivRem</*IsSigned=*/true, /*IsRem=*/false>(div, operands,
+ results);
+ if (auto rem = dyn_cast<arith::RemUIOp>(op))
+ return visitDivRem</*IsSigned=*/false, /*IsRem=*/true>(rem, operands,
+ results);
+ if (auto rem = dyn_cast<arith::RemSIOp>(op))
+ return visitDivRem</*IsSigned=*/true, /*IsRem=*/true>(rem, operands,
+ results);
+ if (auto andi = dyn_cast<arith::AndIOp>(op))
+ return visitAndI(andi, operands, results);
+ if (auto shl = dyn_cast<arith::ShLIOp>(op))
+ return visitShift</*IsLeft=*/true>(shl, operands, results);
+ if (auto shr = dyn_cast<arith::ShRUIOp>(op))
+ return visitShift</*IsLeft=*/false>(shr, operands, results);
+ if (auto sel = dyn_cast<arith::SelectOp>(op))
+ return visitSelect(sel, operands, results);
if (auto cast = dyn_cast<arith::IndexCastOp>(op))
return visitPassThrough(cast, operands, results);
if (auto cast = dyn_cast<arith::IndexCastUIOp>(op))
@@ -230,6 +267,7 @@ class AxisInfoAnalysis
v.contiguity = {n};
v.constancy = {1};
v.divisibility = {kAxisInfoTop};
+ v.innerStride = 1;
propagateIfChanged(results[0], results[0]->join(v));
return success();
}
@@ -308,6 +346,8 @@ class AxisInfoAnalysis
v.constancy[r - 1] = innerConst;
// For a non-AP inner dim, leave at pessimistic.
v.divisibility[r - 1] = baseDiv;
+ if (innerStride != std::numeric_limits<int64_t>::min())
+ v.innerStride = innerStride;
propagateIfChanged(results[0], results[0]->join(v));
return success();
}
@@ -358,6 +398,21 @@ class AxisInfoAnalysis
}
if (src.knownConstant)
v.knownConstant = src.knownConstant;
+ // A broadcast that fans out a scalar / leading-1 source has the broadcast
+ // dim repeating its value -> inner stride 0. Otherwise, the trailing
+ // source dim's stride is preserved when its extent matches the result.
+ auto resShapeArr = resTy.getShape();
+ int64_t innerExt = resShapeArr.back();
+ int sIdxInner = static_cast<int>(rRank - 1) - static_cast<int>(rRank - sRank);
+ if (sIdxInner < 0) {
+ v.innerStride = 0;
+ } else if (srcVt) {
+ int64_t srcInner = srcVt.getShape()[sIdxInner];
+ if (srcInner == 1 && innerExt > 1)
+ v.innerStride = 0;
+ else if (srcInner == innerExt && src.innerStride)
+ v.innerStride = src.innerStride;
+ }
propagateIfChanged(results[0], results[0]->join(v));
return success();
}
@@ -424,6 +479,45 @@ class AxisInfoAnalysis
}
if (src.knownConstant)
v.knownConstant = src.knownConstant;
+ // Identity-like and general row-major reshape both preserve the source
+ // inner-stride property when the source has a single AP characterization.
+ if (src.innerStride)
+ v.innerStride = src.innerStride;
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+
+ // vector.transpose: permute per-dim contiguity / constancy / divisibility
+ // according to the transpose permutation. permutation[i] is the source
+ // dim that ends up at result dim i.
+ LogicalResult visitTranspose(vector::TransposeOp op,
+ ArrayRef<const AxisInfoLattice *> operands,
+ ArrayRef<AxisInfoLattice *> results) {
+ auto resTy = dyn_cast<VectorType>(op.getType());
+ if (!resTy) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ AxisInfo src = operands[0]->getValue();
+ if (!src.isInitialized()) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ ArrayRef<int64_t> perm = op.getPermutation();
+ unsigned r = resTy.getRank();
+ AxisInfo v = AxisInfo::getPessimistic(r);
+ for (unsigned d = 0; d < r; ++d) {
+ unsigned s = static_cast<unsigned>(perm[d]);
+ v.contiguity[d] = src.contiguity[s];
+ v.constancy[d] = src.constancy[s];
+ v.divisibility[d] = src.divisibility[s];
+ }
+ if (src.knownConstant)
+ v.knownConstant = src.knownConstant;
+ // innerStride only survives when the new inner dim came from the old
+ // inner dim (otherwise a different axis is now the contiguous one).
+ if (src.innerStride && perm.back() == src.getRank() - 1)
+ v.innerStride = src.innerStride;
propagateIfChanged(results[0], results[0]->join(v));
return success();
}
@@ -459,6 +553,18 @@ class AxisInfoAnalysis
v.constancy[d] = std::min(lhsConst, rhsConst);
v.divisibility[d] = std::gcd(lhs.divisibility[d], rhs.divisibility[d]);
}
+ // x + uniform-c: stride preserved. x - uniform-c: same. uniform-c - x:
+ // stride flips sign (only useful for the "stride 0" case, which it
+ // preserves trivially).
+ auto isUniform = [&](const AxisInfo &a) {
+ unsigned inner = vt.getRank() - 1;
+ return a.constancy[inner] >= vt.getShape()[inner];
+ };
+ if (lhs.innerStride && isUniform(rhs)) {
+ v.innerStride = *lhs.innerStride;
+ } else if (rhs.innerStride && isUniform(lhs)) {
+ v.innerStride = IsSub ? -*rhs.innerStride : *rhs.innerStride;
+ }
propagateIfChanged(results[0], results[0]->join(v));
return success();
}
@@ -496,6 +602,247 @@ class AxisInfoAnalysis
else
v.contiguity[d] = 1;
}
+ // x * uniform-c: stride scales by c. (Both operands uniform => 0.)
+ unsigned inner = vt.getRank() - 1;
+ auto isUniformInner = [&](const AxisInfo &a) {
+ return a.constancy[inner] >= shape[inner];
+ };
+ if (lhs.innerStride && isUniformInner(rhs) && rhs.knownConstant) {
+ v.innerStride = *lhs.innerStride * *rhs.knownConstant;
+ } else if (rhs.innerStride && isUniformInner(lhs) && lhs.knownConstant) {
+ v.innerStride = *rhs.innerStride * *lhs.knownConstant;
+ }
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+
+ // arith.divui / arith.divsi / arith.remui / arith.remsi by a uniform
+ // positive constant `c`.
+ //
+ // Division (IsRem=false): when the lhs is an inner-dim AP `(base, s, n)`
+ // with `c | s` and `c | divisibility[inner]` (so the per-row base / c is
+ // exact), the result is an AP with stride `s / c` and inner divisibility
+ // `divisibility[inner] / c`. Special cases: `s/c == 1` flags inner-dim
+ // contiguity; `s/c == 0` flags inner-dim constancy.
+ //
+ // Remainder (IsRem=true): when `c | s`, every element of a row sits at
+ // the same residue class -> inner-dim constant -> `innerStride = 0`,
+ // `constancy[inner] = inner` (matches the analysis's notion of
+ // chunk-uniform values along the inner dim).
+ //
+ // Signed vs unsigned only differs in the constant interpretation; we
+ // require positive constants so the signed/unsigned distinction is moot
+ // here.
+ template <bool IsSigned, bool IsRem, typename OpTy>
+ LogicalResult visitDivRem(OpTy op,
+ ArrayRef<const AxisInfoLattice *> operands,
+ ArrayRef<AxisInfoLattice *> results) {
+ auto vt = dyn_cast<VectorType>(op.getType());
+ if (!vt) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ AxisInfo lhs = operands[0]->getValue();
+ AxisInfo rhs = operands[1]->getValue();
+ if (!lhs.isInitialized() || !rhs.isInitialized()) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ unsigned r = vt.getRank();
+ unsigned inner = r - 1;
+ auto shape = vt.getShape();
+ AxisInfo v = AxisInfo::getPessimistic(r);
+
+ bool rhsUniform = rhs.constancy[inner] >= shape[inner] && rhs.knownConstant;
+ if (!rhsUniform || *rhs.knownConstant <= 0) {
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+ int64_t c = *rhs.knownConstant;
+
+ if (!lhs.innerStride) {
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+ int64_t s = *lhs.innerStride;
+ int64_t baseDivLhs = lhs.divisibility[inner];
+ if (s % c != 0) {
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+
+ if (IsRem) {
+ // (base + i*s) mod c, with c | s, is the constant base mod c.
+ v.innerStride = 0;
+ v.constancy[inner] = shape[inner];
+ // The remainder is in [0, c-1], so any power-of-two divisor of c is a
+ // lower bound on alignment. Use lhs's existing divisibility too.
+ v.divisibility[inner] =
+ std::gcd(baseDivLhs, highestPow2Divisor(c));
+ } else {
+ if (baseDivLhs % c != 0) {
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+ int64_t newStride = s / c;
+ v.innerStride = newStride;
+ if (newStride == 1)
+ v.contiguity[inner] = shape[inner];
+ else if (newStride == 0)
+ v.constancy[inner] = shape[inner];
+ v.divisibility[inner] = baseDivLhs / c;
+ }
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+
+ // arith.andi: `x & m` with a uniform positive constant mask `m`.
+ // The most useful case is `x & (P - 1)` for `P` a power of 2: this is
+ // equivalent to `x mod P`, so when the lhs is an inner-dim AP with stride
+ // divisible by `P` the result is inner-dim constant. We also handle the
+ // trivial `m == 0` (always zero) and `m == -1`/all-ones (identity)
+ // shapes.
+ LogicalResult visitAndI(arith::AndIOp op,
+ ArrayRef<const AxisInfoLattice *> operands,
+ ArrayRef<AxisInfoLattice *> results) {
+ auto vt = dyn_cast<VectorType>(op.getType());
+ if (!vt) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ AxisInfo lhs = operands[0]->getValue();
+ AxisInfo rhs = operands[1]->getValue();
+ if (!lhs.isInitialized() || !rhs.isInitialized()) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ unsigned r = vt.getRank();
+ unsigned inner = r - 1;
+ auto shape = vt.getShape();
+ AxisInfo v = AxisInfo::getPessimistic(r);
+
+ // Look for a uniform constant mask on either side.
+ auto getUniformMask = [&](const AxisInfo &a) -> std::optional<int64_t> {
+ if (a.constancy[inner] >= shape[inner] && a.knownConstant)
+ return a.knownConstant;
+ return std::nullopt;
+ };
+ std::optional<int64_t> mLhs = getUniformMask(lhs);
+ std::optional<int64_t> mRhs = getUniformMask(rhs);
+ if (!mLhs && !mRhs) {
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+ const AxisInfo &x = mLhs ? rhs : lhs;
+ int64_t m = mLhs ? *mLhs : *mRhs;
+
+ if (m == 0) {
+ v.knownConstant = 0;
+ v.innerStride = 0;
+ v.constancy[inner] = shape[inner];
+ v.divisibility[inner] = kAxisInfoTop;
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+
+ // `m == P - 1` with P a power of 2 -> equivalent to `x mod P`.
+ if (m > 0 && llvm::isPowerOf2_64(static_cast<uint64_t>(m + 1))) {
+ int64_t P = m + 1;
+ if (x.innerStride && *x.innerStride % P == 0) {
+ v.innerStride = 0;
+ v.constancy[inner] = shape[inner];
+ v.divisibility[inner] =
+ std::gcd(x.divisibility[inner], highestPow2Divisor(P));
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+ }
+ // Conservative fallback for unrecognized masks.
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+
+ // arith.shli (left shift) / arith.shrui (logical right shift) by a
+ // uniform constant `k`. These are `* (1 << k)` and `/ (1 << k)`
+ // (truncating, but for non-negative values the trunc is exact when
+ // `(1 << k)` divides the value). We model them by reducing to mul/divui.
+ template <bool IsLeft, typename OpTy>
+ LogicalResult visitShift(OpTy op,
+ ArrayRef<const AxisInfoLattice *> operands,
+ ArrayRef<AxisInfoLattice *> results) {
+ auto vt = dyn_cast<VectorType>(op.getType());
+ if (!vt) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ AxisInfo lhs = operands[0]->getValue();
+ AxisInfo rhs = operands[1]->getValue();
+ if (!lhs.isInitialized() || !rhs.isInitialized()) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ unsigned r = vt.getRank();
+ unsigned inner = r - 1;
+ auto shape = vt.getShape();
+ AxisInfo v = AxisInfo::getPessimistic(r);
+
+ if (rhs.constancy[inner] < shape[inner] || !rhs.knownConstant) {
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+ int64_t k = *rhs.knownConstant;
+ if (k < 0 || k >= 63) {
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+ int64_t factor = 1LL << k;
+
+ if (IsLeft) {
+ // x << k == x * factor.
+ if (lhs.innerStride) {
+ v.innerStride = *lhs.innerStride * factor;
+ if (*v.innerStride == 1)
+ v.contiguity[inner] = shape[inner];
+ else if (*v.innerStride == 0)
+ v.constancy[inner] = shape[inner];
+ }
+ v.divisibility[inner] = std::min<int64_t>(
+ kAxisInfoTop, lhs.divisibility[inner] * factor);
+ } else {
+ // x >> k == x / factor (for non-negative x); same conditions as divui.
+ if (lhs.innerStride && *lhs.innerStride % factor == 0 &&
+ lhs.divisibility[inner] % factor == 0) {
+ int64_t newStride = *lhs.innerStride / factor;
+ v.innerStride = newStride;
+ if (newStride == 1)
+ v.contiguity[inner] = shape[inner];
+ else if (newStride == 0)
+ v.constancy[inner] = shape[inner];
+ v.divisibility[inner] = lhs.divisibility[inner] / factor;
+ }
+ }
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+
+ // arith.select: result is at least as constrained as the meet of the two
+ // arms. We propagate fields where both arms agree.
+ LogicalResult visitSelect(arith::SelectOp op,
+ ArrayRef<const AxisInfoLattice *> operands,
+ ArrayRef<AxisInfoLattice *> results) {
+ auto vt = dyn_cast<VectorType>(op.getType());
+ if (!vt) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ // operands: [cond, true, false]
+ AxisInfo t = operands[1]->getValue();
+ AxisInfo f = operands[2]->getValue();
+ if (!t.isInitialized() || !f.isInitialized()) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ AxisInfo v = AxisInfo::join(t, f);
propagateIfChanged(results[0], results[0]->join(v));
return success();
}
@@ -527,7 +874,8 @@ using ::mlir::xegpu::detail::axis_dataflow::AxisInfoLattice;
struct CoalesceDecision {
enum class Kind { None, Broadcast, Chunked };
Kind kind = Kind::None;
- int64_t factor = 1; // lane_data factor along the innermost dim
+ int64_t laneLayout = 1; // lane_layout along the innermost dim
+ int64_t factor = 1; // lane_data factor along the innermost dim
};
/// Largest power-of-two `<= bound` that divides `numLanes`.
@@ -547,9 +895,16 @@ static int64_t largestPow2Divisor(int64_t numLanes, int64_t bound) {
}
/// Decide how to coalesce given the offsets axis info.
+///
+/// We pick `lane_layout[inner]` first using the same default rule as
+/// `XeGPUPropagateLayout`: `lane_layout[inner] = min(subgroupSize, inner)`,
+/// rounded down to a power-of-2 divisor of `inner`. The remaining lane
+/// budget then becomes `lane_data[inner] = inner / lane_layout`, capped by
+/// the per-lane chunk-size budget and the offsets contiguity.
static CoalesceDecision decide(const AxisInfo &info,
ArrayRef<int64_t> offsetsShape,
- int64_t origChunk, unsigned maxChunkSize) {
+ int64_t origChunk, unsigned maxChunkSize,
+ unsigned subgroupSize) {
CoalesceDecision d;
if (!info.isInitialized() || offsetsShape.empty())
return d;
@@ -564,18 +919,34 @@ static CoalesceDecision decide(const AxisInfo &info,
return d;
}
+ // Pick lane_layout first: PropagateLayout's default for a 1-D / inner dim
+ // is `min(subgroupSize, inner)`, rounded down to a divisor of inner.
+ int64_t laneLayout = largestPow2Divisor(
+ inner, std::min<int64_t>(subgroupSize, inner));
+ if (laneLayout < 1)
+ laneLayout = 1;
+
+ // Each lane sees `inner / laneLayout` elements of the offsets vector.
+ int64_t perLane = inner / laneLayout;
+ if (perLane < 2)
+ return d; // already one element per lane, nothing to coalesce.
+
if (origChunk < 1)
origChunk = 1;
int64_t budget = static_cast<int64_t>(maxChunkSize) / origChunk;
if (budget < 2)
return d;
+
int64_t bound = std::min<int64_t>(info.contiguity[innerDim], budget);
+ bound = std::min<int64_t>(bound, perLane);
if (bound < 2)
return d;
- int64_t factor = largestPow2Divisor(inner, bound);
+
+ int64_t factor = largestPow2Divisor(perLane, bound);
if (factor < 2)
return d;
d.kind = CoalesceDecision::Kind::Chunked;
+ d.laneLayout = laneLayout;
d.factor = factor;
return d;
}
@@ -594,19 +965,23 @@ static bool isAllTrueMask(Value mask) {
return dense.getSplatValue<APInt>().getBoolValue();
}
-/// Build a `lane_layout`/`lane_data` layout of rank `rank`, with lane_data
-/// = factor on the innermost dim (1 elsewhere) and lane_layout = inner /
-/// factor on the innermost dim (1 elsewhere).
+/// Build a `lane_layout`/`lane_data`/`inst_data` layout of rank `rank`,
+/// with the given lane_layout / lane_data on the innermost dim (1
+/// elsewhere). `inst_data` is `lane_layout * lane_data` per dim, so the
+/// invariant `inst_data[d] == lane_layout[d] * lane_data[d]` holds.
static xegpu::LayoutAttr buildLaneDataLayout(MLIRContext *ctx, unsigned rank,
- int64_t innerLanes,
- int64_t factor) {
+ int64_t innerLaneLayout,
+ int64_t innerLaneData) {
SmallVector<int32_t> laneLayout(rank, 1);
SmallVector<int32_t> laneData(rank, 1);
- laneLayout.back() = static_cast<int32_t>(innerLanes / factor);
- laneData.back() = static_cast<int32_t>(factor);
- return xegpu::LayoutAttr::get(ctx, laneLayout, laneData);
+ SmallVector<int32_t> instData(rank, 1);
+ laneLayout.back() = static_cast<int32_t>(innerLaneLayout);
+ laneData.back() = static_cast<int32_t>(innerLaneData);
+ instData.back() = static_cast<int32_t>(innerLaneLayout * innerLaneData);
+ return xegpu::LayoutAttr::get(ctx, instData, laneLayout, laneData);
}
+
//===----------------------------------------------------------------------===//
// Rewrites.
//===----------------------------------------------------------------------===//
@@ -647,6 +1022,17 @@ static LogicalResult rewriteBroadcastLoad(xegpu::LoadGatherOp op,
namespace {
+/// Look up the subgroup size from the enclosing gpu.module's xevm.target.
+/// Falls back to 16 when no target chip is found or the chip is unknown,
+/// matching the typical Intel Xe2 default. This keeps the pass usable on
+/// plain `module { ... }` IR (e.g. unit lit tests) where there's no
+/// gpu.module / xevm.target wrapper.
+static unsigned lookupSubgroupSize(Operation *op) {
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
+ return uArch ? static_cast<unsigned>(uArch->getSubgroupSize()) : 16u;
+}
+
struct CoalesceLoadPattern final : OpRewritePattern<xegpu::LoadGatherOp> {
CoalesceLoadPattern(MLIRContext *ctx, unsigned maxChunkSize,
DataFlowSolver &solver)
@@ -670,13 +1056,19 @@ struct CoalesceLoadPattern final : OpRewritePattern<xegpu::LoadGatherOp> {
if (!layout.getEffectiveLaneDataAsInt().empty())
return rewriter.notifyMatchFailure(op, "lane_data already set");
+ int64_t origChunk = static_cast<int64_t>(op.getChunkSize().value_or(1));
+ if (op.getChunkSizeAttr() && origChunk > 1)
+ return rewriter.notifyMatchFailure(
+ op, "explicit chunk_size > 1, leaving op alone");
+
+ unsigned subgroupSize = lookupSubgroupSize(op);
+
const auto *lat = solver.lookupState<AxisInfoLattice>(op.getOffsets());
if (!lat || !lat->getValue().isInitialized())
return rewriter.notifyMatchFailure(op, "no axis-info available");
- int64_t origChunk = static_cast<int64_t>(op.getChunkSize().value_or(1));
auto d = decide(lat->getValue(), offsetsTy.getShape(), origChunk,
- maxChunkSize);
+ maxChunkSize, subgroupSize);
if (d.kind == CoalesceDecision::Kind::Broadcast)
return rewriteBroadcastLoad(op, rewriter);
@@ -684,10 +1076,16 @@ struct CoalesceLoadPattern final : OpRewritePattern<xegpu::LoadGatherOp> {
if (d.kind != CoalesceDecision::Kind::Chunked)
return rewriter.notifyMatchFailure(op, "offsets not coalescible");
- int64_t innerLanes = offsetsTy.getShape().back();
auto layout = buildLaneDataLayout(op.getContext(), valueTy.getRank(),
- innerLanes, d.factor);
- rewriter.modifyOpInPlace(op, [&] { op.setLayoutAttr(layout); });
+ d.laneLayout, d.factor);
+ // If the op carried an explicit chunk_size = 1 (the trivial / default
+ // value), drop it: the new lane_data FCD > 1 supersedes it.
+ bool dropChunk = op.getChunkSizeAttr() && origChunk == 1 && d.factor > 1;
+ rewriter.modifyOpInPlace(op, [&] {
+ op.setLayoutAttr(layout);
+ if (dropChunk)
+ op.removeChunkSizeAttr();
+ });
return success();
}
@@ -717,13 +1115,19 @@ struct CoalesceStorePattern final : OpRewritePattern<xegpu::StoreScatterOp> {
if (!layout.getEffectiveLaneDataAsInt().empty())
return rewriter.notifyMatchFailure(op, "lane_data already set");
+ int64_t origChunk = static_cast<int64_t>(op.getChunkSize().value_or(1));
+ if (op.getChunkSizeAttr() && origChunk > 1)
+ return rewriter.notifyMatchFailure(
+ op, "explicit chunk_size > 1, leaving op alone");
+
+ unsigned subgroupSize = lookupSubgroupSize(op);
+
const auto *lat = solver.lookupState<AxisInfoLattice>(op.getOffsets());
if (!lat || !lat->getValue().isInitialized())
return rewriter.notifyMatchFailure(op, "no axis-info available");
- int64_t origChunk = static_cast<int64_t>(op.getChunkSize().value_or(1));
auto d = decide(lat->getValue(), offsetsTy.getShape(), origChunk,
- maxChunkSize);
+ maxChunkSize, subgroupSize);
if (d.kind == CoalesceDecision::Kind::Broadcast)
return rewriter.notifyMatchFailure(
@@ -732,10 +1136,14 @@ struct CoalesceStorePattern final : OpRewritePattern<xegpu::StoreScatterOp> {
if (d.kind != CoalesceDecision::Kind::Chunked)
return rewriter.notifyMatchFailure(op, "offsets not coalescible");
- int64_t innerLanes = offsetsTy.getShape().back();
auto layout = buildLaneDataLayout(op.getContext(), valueTy.getRank(),
- innerLanes, d.factor);
- rewriter.modifyOpInPlace(op, [&] { op.setLayoutAttr(layout); });
+ d.laneLayout, d.factor);
+ bool dropChunk = op.getChunkSizeAttr() && origChunk == 1 && d.factor > 1;
+ rewriter.modifyOpInPlace(op, [&] {
+ op.setLayoutAttr(layout);
+ if (dropChunk)
+ op.removeChunkSizeAttr();
+ });
return success();
}
diff --git a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
index 6f34618f386e4..dd16c1a719fa0 100644
--- a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
+++ b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
@@ -3,11 +3,15 @@
// -----
// vector.step offsets -> stride 1, fully coalescible.
-// 32 lanes, max-chunk-size = 8 -> factor 8 -> lane_layout=[4], lane_data=[8].
+// 32 lanes, subgroup_size = 16 (default) -> lane_layout = 16, perLane = 2,
+// max-chunk-size = 8 -> bound = min(32, 8, 2) = 2 -> lane_data = 2.
+// The trivial `chunk_size = 1` attribute is dropped on success since the
+// new lane_data FCD > 1 supersedes it.
// CHECK-LABEL: func.func @load_step_offsets(
// CHECK: %[[STEP:.*]] = vector.step : vector<32xindex>
// CHECK: %[[LOAD:.*]] = xegpu.load
-// CHECK-SAME: layout = #xegpu.layout<lane_layout = [4], lane_data = [8]>
+// CHECK-NOT: chunk_size
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>
// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
// CHECK: return %[[LOAD]]
func.func @load_step_offsets(%ptr: i64) -> vector<32xf32> {
@@ -19,10 +23,12 @@ func.func @load_step_offsets(%ptr: i64) -> vector<32xf32> {
}
// -----
-// max-chunk-size = 4: factor 4 -> lane_layout=[8], lane_data=[4].
+// max-chunk-size = 4: same lane_layout / lane_data because the bound
+// `min(contiguity = 32, budget = 4, perLane = 2) = 2` already saturates at
+// perLane.
// CHECK4-LABEL: func.func @load_step_offsets_chunk4(
// CHECK4: xegpu.load
-// CHECK4-SAME: layout = #xegpu.layout<lane_layout = [8], lane_data = [4]>
+// CHECK4-SAME: layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>
// CHECK4-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
func.func @load_step_offsets_chunk4(%ptr: i64) -> vector<32xf32> {
%offsets = vector.step : vector<32xindex>
@@ -36,7 +42,7 @@ func.func @load_step_offsets_chunk4(%ptr: i64) -> vector<32xf32> {
// Dense constant arithmetic progression with stride 1.
// CHECK-LABEL: func.func @load_dense_ap_offsets(
// CHECK: xegpu.load
-// CHECK-SAME: layout = #xegpu.layout<lane_layout = [4], lane_data = [8]>
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>
// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xi32>
func.func @load_dense_ap_offsets(%ptr: i64) -> vector<32xi32> {
%offsets = arith.constant dense<[
@@ -99,7 +105,7 @@ func.func @load_partial_mask_unchanged(%ptr: i64, %mask: vector<32xi1>) -> vecto
// Store with vector.step offsets coalesces.
// CHECK-LABEL: func.func @store_step_offsets(
// CHECK: xegpu.store
-// CHECK-SAME: layout = #xegpu.layout<lane_layout = [4], lane_data = [8]>
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>
// CHECK-SAME: : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
func.func @store_step_offsets(%ptr: i64, %v: vector<32xf32>) {
%offsets = vector.step : vector<32xindex>
@@ -127,7 +133,7 @@ func.func @store_broadcast_offsets_unchanged(%ptr: i64, %v: vector<32xf32>) {
// memref-source variant of load coalesces too.
// CHECK-LABEL: func.func @load_memref_step(
// CHECK: xegpu.load
-// CHECK-SAME: layout = #xegpu.layout<lane_layout = [4], lane_data = [8]>
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>
// CHECK-SAME: : memref<1024xf32>, vector<32xindex>, vector<32xi1> -> vector<32xf32>
func.func @load_memref_step(%m: memref<1024xf32>) -> vector<32xf32> {
%offsets = vector.step : vector<32xindex>
@@ -139,10 +145,11 @@ func.func @load_memref_step(%m: memref<1024xf32>) -> vector<32xf32> {
// -----
// 2-D offsets with leading unit dim: inner dim treated as lane dim.
-// vector<1x32xindex> stride-1 -> lane_data=[1, 8], lane_layout=[1, 4].
+// vector<1x32xindex> stride-1, subgroup_size = 16 ->
+// lane_layout = [1, 16], lane_data = [1, 2].
// CHECK-LABEL: func.func @load_2d_leading_unit(
// CHECK: xegpu.load
-// CHECK-SAME: layout = #xegpu.layout<lane_layout = [1, 4], lane_data = [1, 8]>
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>
// CHECK-SAME: : i64, vector<1x32xindex>, vector<1x32xi1> -> vector<1x32xf32>
func.func @load_2d_leading_unit(%ptr: i64) -> vector<1x32xf32> {
%step = vector.step : vector<32xindex>
@@ -154,11 +161,13 @@ func.func @load_2d_leading_unit(%ptr: i64) -> vector<1x32xf32> {
}
// -----
-// True 2-D dense AP: each row stride-1, 16 lanes per row -> lane_data=[1,8],
-// lane_layout=[1, 2].
+// True 2-D dense AP: each row stride-1, inner = 16 = subgroup_size, so
+// perLane = 1 and there's no room for lane_data > 1. The pass leaves the
+// op alone; the wider variant `load_2d_dense_ap_wide` below exercises the
+// coalescing path.
// CHECK-LABEL: func.func @load_2d_dense_ap(
// CHECK: xegpu.load
-// CHECK-SAME: layout = #xegpu.layout<lane_layout = [1, 2], lane_data = [1, 8]>
+// CHECK-NOT: lane_data
// CHECK-SAME: : i64, vector<2x16xindex>, vector<2x16xi1> -> vector<2x16xf32>
func.func @load_2d_dense_ap(%ptr: i64) -> vector<2x16xf32> {
%offsets = arith.constant dense<[
@@ -172,20 +181,47 @@ func.func @load_2d_dense_ap(%ptr: i64) -> vector<2x16xf32> {
}
// -----
-// 1-D step reshape_cast'ed to 4x8: inner extent 8, factor 8
-// -> lane_layout=[1, 1], lane_data=[1, 8] (each of 8 inner positions is a
-// lane that fetches a contiguous chunk of 8).
-// CHECK-LABEL: func.func @load_4x8_step_shape_cast(
+// True 2-D dense AP with inner > subgroup_size: each row stride-1 across 32
+// lanes. With subgroup_size = 16 (from #xevm.target chip = "pvc"),
+// lane_layout[inner] = 16, perLane = 32 / 16 = 2, contiguity[inner] = 32,
+// budget = max-chunk-size / 1 = 8. bound = min(32, 8, 2) = 2 -> lane_data = 2.
+// CHECK-LABEL: gpu.func @load_2d_dense_ap_wide(
// CHECK: xegpu.load
-// CHECK-SAME: layout = #xegpu.layout<lane_layout = [1, 1], lane_data = [1, 8]>
-// CHECK-SAME: : i64, vector<4x8xindex>, vector<4x8xi1> -> vector<4x8xf32>
-func.func @load_4x8_step_shape_cast(%ptr: i64) -> vector<4x8xf32> {
- %step = vector.step : vector<32xindex>
- %offsets = vector.shape_cast %step : vector<32xindex> to vector<4x8xindex>
- %mask = arith.constant dense<true> : vector<4x8xi1>
- %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
- : i64, vector<4x8xindex>, vector<4x8xi1> -> vector<4x8xf32>
- return %v : vector<4x8xf32>
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>
+// CHECK-SAME: : i64, vector<2x32xindex>, vector<2x32xi1> -> vector<2x32xf32>
+gpu.module @kernel [#xevm.target<chip = "pvc">] {
+ gpu.func @load_2d_dense_ap_wide(%ptr: i64) -> vector<2x32xf32> {
+ %offsets = arith.constant dense<[
+ [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
+ 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31],
+ [32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47,
+ 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63]
+ ]> : vector<2x32xindex>
+ %mask = arith.constant dense<true> : vector<2x32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<2x32xindex>, vector<2x32xi1> -> vector<2x32xf32>
+ gpu.return %v : vector<2x32xf32>
+ }
+}
+
+// -----
+// 1-D step reshape_cast'ed to 2x32: inner extent 32 > subgroup_size = 16, so
+// the lane_layout-first rule picks lane_layout[inner] = 16 (perLane = 2),
+// then takes lane_data[inner] = min(contiguity = 32, budget = 8, perLane = 2)
+// rounded down to a power-of-2 divisor of perLane => 2.
+// CHECK-LABEL: gpu.func @load_2x32_step_shape_cast(
+// CHECK: xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>
+// CHECK-SAME: : i64, vector<2x32xindex>, vector<2x32xi1> -> vector<2x32xf32>
+gpu.module @kernel_2x32 [#xevm.target<chip = "pvc">] {
+ gpu.func @load_2x32_step_shape_cast(%ptr: i64) -> vector<2x32xf32> {
+ %step = vector.step : vector<64xindex>
+ %offsets = vector.shape_cast %step : vector<64xindex> to vector<2x32xindex>
+ %mask = arith.constant dense<true> : vector<2x32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<2x32xindex>, vector<2x32xi1> -> vector<2x32xf32>
+ gpu.return %v : vector<2x32xf32>
+ }
}
// -----
@@ -206,11 +242,226 @@ func.func @load_2d_non_ap_unchanged(%ptr: i64) -> vector<2x16xf32> {
return %v : vector<2x16xf32>
}
+// -----
+// "Reduction kernel" pattern with inner > subgroup_size: 2-D offsets built
+// from `transpose(broadcast(rowOffsets))` + `broadcast(step)`. The transpose
+// supplies inner-dim constancy, the broadcast(step) supplies inner-dim
+// contiguity, and the addi recovers contiguity through visitAddSub. With
+// inner = 32 and subgroup_size = 16 this picks lane_layout = [1, 16],
+// lane_data = [1, 2].
+// CHECK-LABEL: gpu.func @load_reduction_pattern_wide(
+// CHECK: xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>
+// CHECK-SAME: : i64, vector<2x32xindex>, vector<2x32xi1> -> vector<2x32xf32>
+gpu.module @kernel_reduction [#xevm.target<chip = "pvc">] {
+ gpu.func @load_reduction_pattern_wide(%ptr: i64, %row0: index, %row1: index)
+ -> vector<2x32xf32> {
+ // Per-row base offsets.
+ %r0 = vector.broadcast %row0 : index to vector<1xindex>
+ %r1 = vector.broadcast %row1 : index to vector<1xindex>
+ %rows = vector.shuffle %r0, %r1 [0, 1] : vector<1xindex>, vector<1xindex>
+ // Inner stride-1 step.
+ %step = vector.step : vector<32xindex>
+ // Build 2x32 offsets: broadcast rows -> transpose -> add broadcast(step).
+ %rowsBc = vector.broadcast %rows
+ : vector<2xindex> to vector<32x2xindex>
+ %rowsT = vector.transpose %rowsBc, [1, 0]
+ : vector<32x2xindex> to vector<2x32xindex>
+ %cols2 = vector.broadcast %step
+ : vector<32xindex> to vector<2x32xindex>
+ %off = arith.addi %rowsT, %cols2 : vector<2x32xindex>
+ %mask = arith.constant dense<true> : vector<2x32xi1>
+ %v = xegpu.load %ptr[%off], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<2x32xindex>, vector<2x32xi1> -> vector<2x32xf32>
+ gpu.return %v : vector<2x32xf32>
+ }
+}
+
+// -----
+// `divui` by a uniform constant equal to the inner stride recovers stride-1
+// contiguity. Here offsets = [0,2,4,…,62] / 2 = [0,1,…,31], 32 lanes
+// against subgroup_size = 16 -> lane_layout = 16, perLane = 2, contiguity
+// = 32, bound = min(32, 8, 2) = 2 -> lane_data = 2.
+// CHECK-LABEL: gpu.func @load_divui_recovers_contiguity(
+// CHECK: xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+gpu.module @kernel_divui [#xevm.target<chip = "pvc">] {
+ gpu.func @load_divui_recovers_contiguity(%ptr: i64) -> vector<32xf32> {
+ %even = arith.constant dense<[
+ 0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30,
+ 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62]>
+ : vector<32xindex>
+ %c2 = arith.constant dense<2> : vector<32xindex>
+ %offsets = arith.divui %even, %c2 : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ gpu.return %v : vector<32xf32>
+ }
+}
+
+// -----
+// `divsi` by a uniform constant equal to the inner stride: same recovery as
+// the divui case.
+// CHECK-LABEL: gpu.func @load_divsi_recovers_contiguity(
+// CHECK: xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+gpu.module @kernel_divsi [#xevm.target<chip = "pvc">] {
+ gpu.func @load_divsi_recovers_contiguity(%ptr: i64) -> vector<32xf32> {
+ %even = arith.constant dense<[
+ 0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30,
+ 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62]>
+ : vector<32xindex>
+ %c2 = arith.constant dense<2> : vector<32xindex>
+ %offsets = arith.divsi %even, %c2 : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ gpu.return %v : vector<32xf32>
+ }
+}
+
+// -----
+// `remui` by a constant that divides the inner stride: every element of a
+// row is the same residue class -> inner-dim constant -> the load becomes a
+// broadcast load (length-1 load + vector.broadcast).
+// CHECK-LABEL: gpu.func @load_remui_inner_uniform(
+// CHECK: %[[L:.*]] = xegpu.load
+// CHECK-SAME: : i64, vector<1xindex>, vector<1xi1> -> vector<1xf32>
+// CHECK: %[[E:.*]] = vector.extract %[[L]][0]
+// CHECK: %[[B:.*]] = vector.broadcast %[[E]] : f32 to vector<16xf32>
+// CHECK: gpu.return %[[B]]
+gpu.module @kernel_remui [#xevm.target<chip = "pvc">] {
+ gpu.func @load_remui_inner_uniform(%ptr: i64) -> vector<16xf32> {
+ %even = arith.constant dense<[
+ 0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30]>
+ : vector<16xindex>
+ %c2 = arith.constant dense<2> : vector<16xindex>
+ // (even % 2) is uniformly 0 along the inner dim.
+ %offsets = arith.remui %even, %c2 : vector<16xindex>
+ %mask = arith.constant dense<true> : vector<16xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ gpu.return %v : vector<16xf32>
+ }
+}
+
+// -----
+// `andi` with a power-of-two-minus-one mask is `% (1 << k)`. With stride 2
+// and mask 1 (= 2-1), every result element is uniform -> broadcast load.
+// CHECK-LABEL: gpu.func @load_andi_inner_uniform(
+// CHECK: %[[L:.*]] = xegpu.load
+// CHECK-SAME: : i64, vector<1xindex>, vector<1xi1> -> vector<1xf32>
+// CHECK: %[[E:.*]] = vector.extract %[[L]][0]
+// CHECK: %[[B:.*]] = vector.broadcast %[[E]] : f32 to vector<16xf32>
+// CHECK: gpu.return %[[B]]
+gpu.module @kernel_andi [#xevm.target<chip = "pvc">] {
+ gpu.func @load_andi_inner_uniform(%ptr: i64) -> vector<16xf32> {
+ %even = arith.constant dense<[
+ 0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30]>
+ : vector<16xindex>
+ %m = arith.constant dense<1> : vector<16xindex>
+ %offsets = arith.andi %even, %m : vector<16xindex>
+ %mask = arith.constant dense<true> : vector<16xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ gpu.return %v : vector<16xf32>
+ }
+}
+
+// -----
+// `shli` by a constant scales the inner stride. step << 1 -> stride 2,
+// which on its own is not coalescible (no divui follows), so we get no
+// layout attached even when there's room (vector<32>, perLane=2).
+// CHECK-LABEL: gpu.func @load_shli_unchanged(
+// CHECK: xegpu.load
+// CHECK-NOT: lane_data
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+gpu.module @kernel_shli [#xevm.target<chip = "pvc">] {
+ gpu.func @load_shli_unchanged(%ptr: i64) -> vector<32xf32> {
+ %step = vector.step : vector<32xindex>
+ %k = arith.constant dense<1> : vector<32xindex>
+ %offsets = arith.shli %step, %k : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ gpu.return %v : vector<32xf32>
+ }
+}
+
+// -----
+// `shli` followed by `shrui` cancels: `(step << 1) >> 1` has innerStride
+// scaled to 2 then divided by 2 -> stride 1 again. With inner = 32 and
+// subgroup_size = 16 we coalesce by lane_data = 2.
+// CHECK-LABEL: gpu.func @load_shli_then_shrui_recovers(
+// CHECK: xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+gpu.module @kernel_shli_shrui [#xevm.target<chip = "pvc">] {
+ gpu.func @load_shli_then_shrui_recovers(%ptr: i64) -> vector<32xf32> {
+ %step = vector.step : vector<32xindex>
+ %k = arith.constant dense<1> : vector<32xindex>
+ %doubled = arith.shli %step, %k : vector<32xindex>
+ %offsets = arith.shrui %doubled, %k : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ gpu.return %v : vector<32xf32>
+ }
+}
+
+// -----
+// `arith.select` between two AP arms with the same inner-dim properties:
+// the result inherits the meet of the two arms. Both arms here are
+// stride-1 step + a per-arm constant base, so the select preserves
+// inner-dim contiguity = 32 -> coalesces with lane_data = 2.
+// CHECK-LABEL: gpu.func @load_select_two_aps(
+// CHECK: xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+gpu.module @kernel_select [#xevm.target<chip = "pvc">] {
+ gpu.func @load_select_two_aps(%ptr: i64, %cond: i1) -> vector<32xf32> {
+ %step = vector.step : vector<32xindex>
+ %a = arith.constant dense<0> : vector<32xindex>
+ %b = arith.constant dense<64> : vector<32xindex>
+ %baseA = arith.addi %step, %a : vector<32xindex>
+ %baseB = arith.addi %step, %b : vector<32xindex>
+ %offsets = arith.select %cond, %baseA, %baseB : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ gpu.return %v : vector<32xf32>
+ }
+}
+
+// -----
+// `divui` by a constant that does NOT divide the inner stride: stride 2 / 3
+// is not exact, so we conservatively give up. No layout attached.
+// CHECK-LABEL: gpu.func @load_divui_non_divisor_unchanged(
+// CHECK: xegpu.load
+// CHECK-NOT: lane_data
+// CHECK-SAME: : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+gpu.module @kernel_divui_neg [#xevm.target<chip = "pvc">] {
+ gpu.func @load_divui_non_divisor_unchanged(%ptr: i64) -> vector<16xf32> {
+ %even = arith.constant dense<[
+ 0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30]>
+ : vector<16xindex>
+ %c3 = arith.constant dense<3> : vector<16xindex>
+ %offsets = arith.divui %even, %c3 : vector<16xindex>
+ %mask = arith.constant dense<true> : vector<16xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ gpu.return %v : vector<16xf32>
+ }
+}
+
// -----
// 2-D store with step + shape_cast offsets coalesces too.
// CHECK-LABEL: func.func @store_2d_step(
// CHECK: xegpu.store
-// CHECK-SAME: layout = #xegpu.layout<lane_layout = [1, 4], lane_data = [1, 8]>
+// CHECK-SAME: layout = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>
// CHECK-SAME: : vector<1x32xf32>, i64, vector<1x32xindex>, vector<1x32xi1>
func.func @store_2d_step(%ptr: i64, %v: vector<1x32xf32>) {
%step = vector.step : vector<32xindex>
@@ -220,3 +471,24 @@ func.func @store_2d_step(%ptr: i64, %v: vector<1x32xf32>) {
: vector<1x32xf32>, i64, vector<1x32xindex>, vector<1x32xi1>
return
}
+
+// -----
+// An op that already declares chunk_size = 2 is left alone: the verifier
+// requires a particular value/mask shape relationship for chunked ops, so
+// the pass conservatively skips when an explicit chunk_size > 1 is set
+// (a downstream pass has already committed to that per-lane chunked
+// access).
+// CHECK-LABEL: gpu.func @load_explicit_chunk_unchanged(
+// CHECK: xegpu.load
+// CHECK-NOT: lane_data
+// CHECK-SAME: <{chunk_size = 2 : i64}>
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32x2xf32>
+gpu.module @kernel_explicit_chunk [#xevm.target<chip = "pvc">] {
+ gpu.func @load_explicit_chunk_unchanged(%ptr: i64) -> vector<32x2xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 2 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32x2xf32>
+ gpu.return %v : vector<32x2xf32>
+ }
+}
>From 2d788139042519c3726927a7e06cb9e5239ac7cc Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Thu, 28 May 2026 22:17:51 +0000
Subject: [PATCH 4/5] [mlir][XeGPU][Transform] Disable broadcast-load coalesce
path; drop the corresponding lit test.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Gate `rewriteBroadcastLoad` off behind a TODO in `CoalesceLoadPattern` —
the Broadcast decision still falls through to "not coalescible" and the
op is left unchanged. The store-side broadcast skip is unaffected.
Also tidy the file with clang-format-style reflows on a few long lines
(no semantic change), and update the lit test:
- Remove `load_broadcast_offsets` (asserted the rewrite that's now
disabled).
- Convert `load_remui_inner_uniform` and `load_andi_inner_uniform` from
asserting a broadcast-load output to asserting the load is left alone
(the analysis still classifies them as Broadcast; the rewrite just
doesn't fire).
Co-Authored-By: Claude Opus 4.7 <noreply at anthropic.com>
---
.../Transforms/XeGPUCoalesceGatherScatter.cpp | 42 ++++++++-----------
.../XeGPU/coalesce-gather-scatter.mlir | 40 +++++-------------
2 files changed, 29 insertions(+), 53 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
index 4320ccb2aa6bb..99cac5a724b6b 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
@@ -320,8 +320,7 @@ class AxisInfoAnalysis
auto values = llvm::to_vector(dense.getValues<APInt>());
int64_t innerCont = inner;
int64_t innerConst = inner;
- int64_t innerStride =
- values[1].getSExtValue() - values[0].getSExtValue();
+ int64_t innerStride = values[1].getSExtValue() - values[0].getSExtValue();
int64_t base = values[0].getSExtValue();
int64_t baseDiv = highestPow2Divisor(base);
for (int64_t o = 0; o < outer; ++o) {
@@ -380,16 +379,14 @@ class AxisInfoAnalysis
if (sIdx < 0) {
v.constancy[d] = resExt;
v.contiguity[d] = 1;
- v.divisibility[d] =
- src.isInitialized() ? src.divisibility.front() : 1;
+ v.divisibility[d] = src.isInitialized() ? src.divisibility.front() : 1;
continue;
}
int64_t srcExt = srcVt.getShape()[sIdx];
if (srcExt == 1 && resExt > 1) {
v.constancy[d] = resExt;
v.contiguity[d] = 1;
- v.divisibility[d] =
- src.isInitialized() ? src.divisibility[sIdx] : 1;
+ v.divisibility[d] = src.isInitialized() ? src.divisibility[sIdx] : 1;
} else if (src.isInitialized()) {
v.contiguity[d] = src.contiguity[sIdx];
v.constancy[d] = src.constancy[sIdx];
@@ -403,7 +400,8 @@ class AxisInfoAnalysis
// source dim's stride is preserved when its extent matches the result.
auto resShapeArr = resTy.getShape();
int64_t innerExt = resShapeArr.back();
- int sIdxInner = static_cast<int>(rRank - 1) - static_cast<int>(rRank - sRank);
+ int sIdxInner =
+ static_cast<int>(rRank - 1) - static_cast<int>(rRank - sRank);
if (sIdxInner < 0) {
v.innerStride = 0;
} else if (srcVt) {
@@ -523,8 +521,7 @@ class AxisInfoAnalysis
}
template <bool IsSub, typename OpTy>
- LogicalResult visitAddSub(OpTy op,
- ArrayRef<const AxisInfoLattice *> operands,
+ LogicalResult visitAddSub(OpTy op, ArrayRef<const AxisInfoLattice *> operands,
ArrayRef<AxisInfoLattice *> results) {
auto vt = dyn_cast<VectorType>(op.getType());
if (!vt) {
@@ -547,8 +544,8 @@ class AxisInfoAnalysis
// contiguity propagates through add when one side is constant on the
// run, and through sub only when the rhs is constant on the run.
int64_t cont = IsSub ? std::min(lhsCont, rhsConst)
- : std::max(std::min(lhsCont, rhsConst),
- std::min(rhsCont, lhsConst));
+ : std::max(std::min(lhsCont, rhsConst),
+ std::min(rhsCont, lhsConst));
v.contiguity[d] = std::max<int64_t>(1, cont);
v.constancy[d] = std::min(lhsConst, rhsConst);
v.divisibility[d] = std::gcd(lhs.divisibility[d], rhs.divisibility[d]);
@@ -634,8 +631,7 @@ class AxisInfoAnalysis
// require positive constants so the signed/unsigned distinction is moot
// here.
template <bool IsSigned, bool IsRem, typename OpTy>
- LogicalResult visitDivRem(OpTy op,
- ArrayRef<const AxisInfoLattice *> operands,
+ LogicalResult visitDivRem(OpTy op, ArrayRef<const AxisInfoLattice *> operands,
ArrayRef<AxisInfoLattice *> results) {
auto vt = dyn_cast<VectorType>(op.getType());
if (!vt) {
@@ -677,8 +673,7 @@ class AxisInfoAnalysis
v.constancy[inner] = shape[inner];
// The remainder is in [0, c-1], so any power-of-two divisor of c is a
// lower bound on alignment. Use lhs's existing divisibility too.
- v.divisibility[inner] =
- std::gcd(baseDivLhs, highestPow2Divisor(c));
+ v.divisibility[inner] = std::gcd(baseDivLhs, highestPow2Divisor(c));
} else {
if (baseDivLhs % c != 0) {
propagateIfChanged(results[0], results[0]->join(v));
@@ -767,8 +762,7 @@ class AxisInfoAnalysis
// (truncating, but for non-negative values the trunc is exact when
// `(1 << k)` divides the value). We model them by reducing to mul/divui.
template <bool IsLeft, typename OpTy>
- LogicalResult visitShift(OpTy op,
- ArrayRef<const AxisInfoLattice *> operands,
+ LogicalResult visitShift(OpTy op, ArrayRef<const AxisInfoLattice *> operands,
ArrayRef<AxisInfoLattice *> results) {
auto vt = dyn_cast<VectorType>(op.getType());
if (!vt) {
@@ -806,8 +800,8 @@ class AxisInfoAnalysis
else if (*v.innerStride == 0)
v.constancy[inner] = shape[inner];
}
- v.divisibility[inner] = std::min<int64_t>(
- kAxisInfoTop, lhs.divisibility[inner] * factor);
+ v.divisibility[inner] =
+ std::min<int64_t>(kAxisInfoTop, lhs.divisibility[inner] * factor);
} else {
// x >> k == x / factor (for non-negative x); same conditions as divui.
if (lhs.innerStride && *lhs.innerStride % factor == 0 &&
@@ -921,8 +915,8 @@ static CoalesceDecision decide(const AxisInfo &info,
// Pick lane_layout first: PropagateLayout's default for a 1-D / inner dim
// is `min(subgroupSize, inner)`, rounded down to a divisor of inner.
- int64_t laneLayout = largestPow2Divisor(
- inner, std::min<int64_t>(subgroupSize, inner));
+ int64_t laneLayout =
+ largestPow2Divisor(inner, std::min<int64_t>(subgroupSize, inner));
if (laneLayout < 1)
laneLayout = 1;
@@ -981,7 +975,6 @@ static xegpu::LayoutAttr buildLaneDataLayout(MLIRContext *ctx, unsigned rank,
return xegpu::LayoutAttr::get(ctx, instData, laneLayout, laneData);
}
-
//===----------------------------------------------------------------------===//
// Rewrites.
//===----------------------------------------------------------------------===//
@@ -1070,8 +1063,9 @@ struct CoalesceLoadPattern final : OpRewritePattern<xegpu::LoadGatherOp> {
auto d = decide(lat->getValue(), offsetsTy.getShape(), origChunk,
maxChunkSize, subgroupSize);
- if (d.kind == CoalesceDecision::Kind::Broadcast)
- return rewriteBroadcastLoad(op, rewriter);
+ // TODO: Don't do broadcast coalescing for now.
+ // if (d.kind == CoalesceDecision::Kind::Broadcast)
+ // return rewriteBroadcastLoad(op, rewriter);
if (d.kind != CoalesceDecision::Kind::Chunked)
return rewriter.notifyMatchFailure(op, "offsets not coalescible");
diff --git a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
index dd16c1a719fa0..859580c95906f 100644
--- a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
+++ b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
@@ -55,22 +55,6 @@ func.func @load_dense_ap_offsets(%ptr: i64) -> vector<32xi32> {
return %v : vector<32xi32>
}
-// -----
-// All-equal offsets -> broadcast load (single load + vector.broadcast).
-// CHECK-LABEL: func.func @load_broadcast_offsets(
-// CHECK: %[[L:.*]] = xegpu.load
-// CHECK-SAME: : i64, vector<1xindex>, vector<1xi1> -> vector<1xf32>
-// CHECK: %[[E:.*]] = vector.extract %[[L]][0]
-// CHECK: %[[B:.*]] = vector.broadcast %[[E]] : f32 to vector<32xf32>
-// CHECK: return %[[B]]
-func.func @load_broadcast_offsets(%ptr: i64) -> vector<32xf32> {
- %offsets = arith.constant dense<0> : vector<32xindex>
- %mask = arith.constant dense<true> : vector<32xi1>
- %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
- : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- return %v : vector<32xf32>
-}
-
// -----
// Non-stride-1 (stride 4) offsets: not contiguous, no layout attached.
// CHECK-LABEL: func.func @load_stride4_unchanged(
@@ -325,14 +309,13 @@ gpu.module @kernel_divsi [#xevm.target<chip = "pvc">] {
// -----
// `remui` by a constant that divides the inner stride: every element of a
-// row is the same residue class -> inner-dim constant -> the load becomes a
-// broadcast load (length-1 load + vector.broadcast).
+// row is the same residue class -> inner-dim constant. The decision picks
+// the broadcast case, which is currently disabled, so the load is left
+// alone (no layout, no chunk_size change).
// CHECK-LABEL: gpu.func @load_remui_inner_uniform(
-// CHECK: %[[L:.*]] = xegpu.load
-// CHECK-SAME: : i64, vector<1xindex>, vector<1xi1> -> vector<1xf32>
-// CHECK: %[[E:.*]] = vector.extract %[[L]][0]
-// CHECK: %[[B:.*]] = vector.broadcast %[[E]] : f32 to vector<16xf32>
-// CHECK: gpu.return %[[B]]
+// CHECK: xegpu.load
+// CHECK-NOT: lane_data
+// CHECK-SAME: : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
gpu.module @kernel_remui [#xevm.target<chip = "pvc">] {
gpu.func @load_remui_inner_uniform(%ptr: i64) -> vector<16xf32> {
%even = arith.constant dense<[
@@ -350,13 +333,12 @@ gpu.module @kernel_remui [#xevm.target<chip = "pvc">] {
// -----
// `andi` with a power-of-two-minus-one mask is `% (1 << k)`. With stride 2
-// and mask 1 (= 2-1), every result element is uniform -> broadcast load.
+// and mask 1 (= 2-1), every result element is uniform: same broadcast case
+// as remui above, currently disabled, so the load is left alone.
// CHECK-LABEL: gpu.func @load_andi_inner_uniform(
-// CHECK: %[[L:.*]] = xegpu.load
-// CHECK-SAME: : i64, vector<1xindex>, vector<1xi1> -> vector<1xf32>
-// CHECK: %[[E:.*]] = vector.extract %[[L]][0]
-// CHECK: %[[B:.*]] = vector.broadcast %[[E]] : f32 to vector<16xf32>
-// CHECK: gpu.return %[[B]]
+// CHECK: xegpu.load
+// CHECK-NOT: lane_data
+// CHECK-SAME: : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
gpu.module @kernel_andi [#xevm.target<chip = "pvc">] {
gpu.func @load_andi_inner_uniform(%ptr: i64) -> vector<16xf32> {
%even = arith.constant dense<[
>From 459c1677ea7a9067b2d95be8b11799d8d8978329 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Mon, 1 Jun 2026 23:50:10 +0000
Subject: [PATCH 5/5] [mlir][XeGPU][Transform][Test] Split
xegpu-coalesce-gather-scatter into analysis + apply APIs; move pass to
test-only; remove broadcast-load rewrite.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Restructure the coalesce-gather-scatter functionality so another
pass (e.g. xegpu-propagate-layout) can take over the per-op decision
of whether to honor a coalescing hint.
- New `XeGPU_CoalesceHintAttr` (`#xegpu.coalesce_hint<factor = N>`):
a discardable attribute carrying the chosen FCD `lane_data` factor.
Only `factor` is stored; `lane_layout[FCD]` is re-derived at apply
time from the op's offsets inner extent and the chip-derived
subgroup size. Verifier requires `factor` to be a power of two
>= 2.
- New public APIs in `mlir/Dialect/XeGPU/Transforms/Transforms.h`:
- `runCoalesceGatherScatterAnalysis(root, opts)`: runs the
AxisInfo-based offset analysis and stamps a
`xegpu.coalesce_hint` attribute on every coalescible
`xegpu.load` / `xegpu.store`. No rewrite.
- `applyCoalesceGatherScatterHint(op)` /
`applyCoalesceGatherScatterHints(root)`: read a stamped hint
and install the equivalent layout, drop a trivial
`chunk_size = 1`, remove the hint. Idempotent.
- `clearCoalesceGatherScatterHints(root)`: strip leftover
hints (for callers that decided not to apply some hints).
- Remove the broadcast-load rewrite (`rewriteBroadcastLoad` and the
`Broadcast` decision kind). The analysis still classifies
inner-dim-uniform offsets internally but reports them as
"not coalescible" — the previous length-1-load + `vector.broadcast`
emission conflicts with downstream layout propagation, and a
proper rewrite belongs in a dedicated pass when needed.
- Remove the standalone `xegpu-coalesce-gather-scatter` pass:
- Drop the entry from `Passes.td`.
- Drop the call site from `GPUToXeVMPipeline.cpp`. Coalescing is
no longer scheduled in the production pipeline; integration
into `xegpu-propagate-layout` is being implemented on a
separate branch.
- Add `TestXeGPUCoalesceGatherScatter` in
`mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp`, registered as
`test-xegpu-coalesce-gather-scatter`. The pass body is a thin
driver that calls the analysis API and (unless `analyze-only=true`)
the apply API. `OperationPass<>` so existing `func.func`-rooted
tests still match.
Lit:
- `mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir` retargeted
to the new CLI flag. Same coverage as before.
- New `mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir`
exercises the analysis API in isolation (`analyze-only=true`),
pinning the `xegpu.coalesce_hint` attribute contract that the
upcoming propagate-layout integration will consume.
Co-Authored-By: Claude Opus 4.7 <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/IR/XeGPUAttrs.td | 41 +++
.../mlir/Dialect/XeGPU/Transforms/Passes.td | 47 ---
.../Dialect/XeGPU/Transforms/Transforms.h | 39 +++
.../GPU/Pipelines/GPUToXeVMPipeline.cpp | 2 -
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 15 +
.../Transforms/XeGPUCoalesceGatherScatter.cpp | 320 ++++++++----------
.../coalesce-gather-scatter-analyze.mlir | 105 ++++++
.../XeGPU/coalesce-gather-scatter.mlir | 4 +-
.../lib/Dialect/XeGPU/TestXeGPUTransforms.cpp | 43 +++
9 files changed, 377 insertions(+), 239 deletions(-)
create mode 100644 mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
index 40edce8a60429..50b67cefafa31 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
@@ -927,6 +927,47 @@ def XeGPU_MemLayoutAttr : XeGPUAttr<"MemLayout", "mem_layout"> {
}
+def XeGPU_CoalesceHintAttr : XeGPUAttr<"CoalesceHint", "coalesce_hint"> {
+ let summary = [{Per-op hint stamped by the coalesce-gather-scatter analysis.}];
+
+ let description = [{
+ `CoalesceHintAttr` is a discardable attribute attached by the
+ coalesce-gather-scatter analysis to `xegpu.load` / `xegpu.store` ops
+ whose offsets describe a contiguous-per-lane access pattern. It records
+ the chosen `lane_data` factor along the fastest-changing dim (FCD) of
+ the value vector. The `lane_layout` along the FCD is *not* stored here;
+ it is re-derived at apply time from the op's offsets inner extent and
+ the chip's subgroup size, mirroring `xegpu-propagate-layout`'s default
+ rule.
+
+ The attribute is consumed and removed by `applyCoalesceGatherScatterHint`,
+ which installs an equivalent `xegpu.layout` on the op. A producer that
+ decides not to apply the hint (for example, a propagator that detects a
+ conflict with an anchor-driven layout) should remove the attribute
+ rather than leave it dangling.
+
+ Example:
+ ```mlir
+ %v = xegpu.load %ptr[%offsets], %mask
+ <{chunk_size = 1 : i64,
+ xegpu.coalesce_hint = #xegpu.coalesce_hint<factor = 4>}>
+ : i64, vector<64xindex>, vector<64xi1> -> vector<64xf32>
+ ```
+ }];
+
+ let parameters = (ins "IntegerAttr": $factor);
+
+ let builders = [
+ AttrBuilder<(ins "int64_t":$factor), [{
+ return $_get($_ctxt, IntegerAttr::get(IntegerType::get($_ctxt, 64),
+ factor));
+ }]>
+ ];
+
+ let assemblyFormat = "`<` `factor` `=` $factor `>`";
+ let genVerifyDecl = 1;
+}
+
def AnchorLayoutInterface : OpInterface<"AnchorLayoutInterface"> {
let cppNamespace = "::mlir::xegpu";
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
index caf33cc1e13c9..227d36653eb9d 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
@@ -118,51 +118,4 @@ def XeGPUSgToWiDistributeExperimental : Pass<"xegpu-sg-to-wi-distribute-experime
"vector::VectorDialect", "index::IndexDialect"];
}
-def XeGPUCoalesceGatherScatter : Pass<"xegpu-coalesce-gather-scatter"> {
- let summary = "Coalesce neighbouring lanes of xegpu.load/store into chunked accesses";
- let description = [{
- Rewrites `xegpu.load` / `xegpu.store` ops whose offsets vector describes a
- contiguous-per-lane access into an equivalent op with a smaller offsets
- vector and a larger `chunk_size`. The transformation reduces the number of
- memory messages by having each lane fetch / write multiple contiguous
- elements.
-
- Two index patterns are recognized:
- - Affine offsets of the form `base + i * stride` (e.g. `vector.step` or a
- dense constant arithmetic progression). The offsets are coalescible with
- a factor `N` when `stride` divides `N` and the original `chunk_size`
- multiplied by `N` does not exceed the configured maximum.
- - All-equal offsets (e.g. `dense<0>`). All lanes load from the same
- address; this is rewritten as a broadcast: a single load of one element
- is materialized and the result is splatted via `vector.broadcast`.
-
- The mask must also be uniformly `true` (a `dense<true>` constant) for the
- coalesced lanes; partial-mask coalescing is not yet supported.
-
- This pass is intended to run before `xegpu-propagate-layout` so that the
- coalesced shape participates in `inst_data` selection.
-
- Pass options:
- - `max-chunk-size`: upper bound on the produced `chunk_size`. Defaults to
- 8, matching typical Xe scatter chunk-size limits.
-
- Ops that already declare an explicit `chunk_size > 1` are left alone:
- a downstream pass has already committed to a per-lane chunked access,
- and the verifier-imposed shape relationship between value, offsets,
- and mask makes coalescing unsafe without a deeper rewrite.
-
- The subgroup size is read per-op from the enclosing `gpu.module`'s
- `xevm.target` chip (the same lookup as `xegpu-propagate-layout`). If no
- target chip is found, the pass falls back to a subgroup size of 16.
- }];
- let dependentDialects = [
- "arith::ArithDialect", "memref::MemRefDialect", "xegpu::XeGPUDialect",
- "vector::VectorDialect"];
- let options = [Option<
- "maxChunkSize", "max-chunk-size", "unsigned",
- /*default=*/"8",
- "Upper bound on the produced chunk_size when coalescing.">];
-}
-
-
#endif // MLIR_DIALECT_XEGPU_TRANSFORMS_PASSES_TD
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
index a21866b5cc33f..eeaf0a8809ea7 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
@@ -83,6 +83,45 @@ void populateXeGPUSgToWiDistributeTypeConversionAndLegality(
TypeConverter &typeConverter, RewritePatternSet &patterns,
ConversionTarget &target);
+//===----------------------------------------------------------------------===//
+// Coalesce gather/scatter analysis + apply.
+//===----------------------------------------------------------------------===//
+
+/// Options controlling `runCoalesceGatherScatterAnalysis`.
+struct CoalesceGatherScatterAnalysisOptions {
+ /// Upper bound on the per-lane chunk size produced by coalescing. Mirrors
+ /// the `max-chunk-size` option of the original pass.
+ unsigned maxChunkSize = 8;
+};
+
+/// Run the AxisInfo-based coalescing analysis over `root` and stamp a
+/// `xegpu.coalesce_hint` attribute on every `xegpu.load` / `xegpu.store`
+/// the analysis classifies as coalescible. Ops with an existing non-empty
+/// `lane_data`, an explicit `chunk_size > 1`, or a non-uniform mask are
+/// skipped (no hint stamped).
+///
+/// This function performs no rewrite of its own; the hint is consumed by
+/// `applyCoalesceGatherScatterHint` (or a downstream pass that wants
+/// stronger control over when to honor the hint).
+void runCoalesceGatherScatterAnalysis(
+ Operation *root, const CoalesceGatherScatterAnalysisOptions &options = {});
+
+/// Apply a stamped `xegpu.coalesce_hint` on `op`: install an equivalent
+/// `lane_layout` / `lane_data` / `inst_data` layout, drop a trivial
+/// `chunk_size = 1` attribute, and remove the hint. Idempotent — no-op if
+/// the op carries no hint. Returns `failure()` when the hint is malformed
+/// (e.g. attached to an op that isn't a gather/scatter).
+LogicalResult applyCoalesceGatherScatterHint(Operation *op);
+
+/// Walk `root` and apply coalesce hints on every op that carries one.
+/// Hints stamped on unsupported ops are silently dropped.
+void applyCoalesceGatherScatterHints(Operation *root);
+
+/// Walk `root` and remove any leftover `xegpu.coalesce_hint` attributes —
+/// useful as a cleanup after a propagator has decided whether to honor each
+/// hint.
+void clearCoalesceGatherScatterHints(Operation *root);
+
/// Collect a set of patterns to unroll xegpu operations to a smaller shapes.
/// Users can control whether an operation to be unrolled or not, as well as
/// its target shape via `options` structure. (via setting filterConstraint
diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
index c515721363a73..7600ec39fb3f5 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
@@ -74,8 +74,6 @@ void buildGPUPassPipeline(OpPassManager &pm,
pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
pm.addNestedPass<gpu::GPUModuleOp>(createLowerAffinePass());
pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
- pm.addNestedPass<gpu::GPUModuleOp>(
- xegpu::createXeGPUCoalesceGatherScatter());
xegpu::XeGPUPropagateLayoutOptions instDataOptions;
instDataOptions.layoutKind = "inst";
pm.addNestedPass<gpu::GPUModuleOp>(
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index e92b109c2223e..8cdef2dd994d8 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -1213,6 +1213,21 @@ RangeAttr::verify(llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
return success();
}
+//===----------------------------------------------------------------------===//
+// XeGPU_CoalesceHintAttr
+//===----------------------------------------------------------------------===//
+
+LogicalResult CoalesceHintAttr::verify(
+ llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
+ IntegerAttr factor) {
+ int64_t f = factor.getInt();
+ if (f < 2)
+ return emitError() << "'factor' : " << f << " must be >= 2";
+ if ((f & (f - 1)) != 0)
+ return emitError() << "'factor' : " << f << " must be a power of two";
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// XeGPU_TensorDescType
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
index 99cac5a724b6b..c8ab290b9d654 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
@@ -21,11 +21,11 @@
// (e.g. `vector<1x32xindex>`) are handled by treating the inner dim as the
// lane dim.
//
-// The `Broadcast` case (constancy along the innermost dim equals the inner
-// length) is special: there is no layout-only encoding for "all lanes load
-// the same scalar", so for loads we still rewrite to a length-1
-// `xegpu.load` followed by `vector.broadcast`. Stores in this shape are
-// skipped (last-writer-wins is ambiguous).
+// All-equal offsets ("uniform inner dim") are detected by the analysis
+// but the pass currently leaves such ops alone — there is no layout-only
+// encoding for "all lanes load the same scalar", and the previous
+// length-1-load + `vector.broadcast` rewrite was removed because it
+// conflicts with downstream layout propagation.
//
//===----------------------------------------------------------------------===//
@@ -37,6 +37,7 @@
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
+#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
#include "mlir/IR/BuiltinAttributes.h"
@@ -50,13 +51,6 @@
#include <numeric>
#include <optional>
-namespace mlir {
-namespace xegpu {
-#define GEN_PASS_DEF_XEGPUCOALESCEGATHERSCATTER
-#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"
-} // namespace xegpu
-} // namespace mlir
-
#define DEBUG_TYPE "xegpu-coalesce-gather-scatter"
using namespace mlir;
@@ -866,7 +860,7 @@ using ::mlir::xegpu::detail::axis_dataflow::AxisInfoLattice;
//===----------------------------------------------------------------------===//
struct CoalesceDecision {
- enum class Kind { None, Broadcast, Chunked };
+ enum class Kind { None, Chunked };
Kind kind = Kind::None;
int64_t laneLayout = 1; // lane_layout along the innermost dim
int64_t factor = 1; // lane_data factor along the innermost dim
@@ -907,11 +901,10 @@ static CoalesceDecision decide(const AxisInfo &info,
if (inner < 2)
return d;
- // Broadcast if the innermost dim is uniform across all lanes.
- if (info.constancy[innerDim] >= inner) {
- d.kind = CoalesceDecision::Kind::Broadcast;
+ // All-equal offsets: every lane sees the same address. There is no
+ // layout-only encoding for this; we leave the op alone.
+ if (info.constancy[innerDim] >= inner)
return d;
- }
// Pick lane_layout first: PropagateLayout's default for a 1-D / inner dim
// is `min(subgroupSize, inner)`, rounded down to a divisor of inner.
@@ -979,40 +972,6 @@ static xegpu::LayoutAttr buildLaneDataLayout(MLIRContext *ctx, unsigned rank,
// Rewrites.
//===----------------------------------------------------------------------===//
-/// Replace an `xegpu.load` whose offsets are uniform along the innermost
-/// dim with a length-1 load + `vector.broadcast` back to the original
-/// value type. Works for any rank; the length-1 load uses an inner-dim
-/// length-1 offsets/mask vector.
-static LogicalResult rewriteBroadcastLoad(xegpu::LoadGatherOp op,
- PatternRewriter &rewriter) {
- Location loc = op.getLoc();
- auto valueTy = op.getValueType();
- auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
- if (!valueTy || !offsetsTy)
- return failure();
-
- // Extract a scalar offset from index 0...0.
- SmallVector<int64_t> zeros(offsetsTy.getRank(), 0);
- Value scalarOffset =
- vector::ExtractOp::create(rewriter, loc, op.getOffsets(), zeros);
- auto idxVecTy = VectorType::get({1}, rewriter.getIndexType());
- auto maskVecTy = VectorType::get({1}, rewriter.getI1Type());
- Value newOffsets =
- vector::BroadcastOp::create(rewriter, loc, idxVecTy, scalarOffset);
- Value newMask = arith::ConstantOp::create(
- rewriter, loc, DenseIntElementsAttr::get(maskVecTy, true));
- auto newValueTy = VectorType::get({1}, valueTy.getElementType());
- auto newLoad = xegpu::LoadGatherOp::create(
- rewriter, loc, newValueTy, op.getSource(), newOffsets, newMask,
- /*chunk_size=*/IntegerAttr(), op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), /*layout=*/xegpu::DistributeLayoutAttr());
- Value scalar = vector::ExtractOp::create(rewriter, loc, newLoad.getResult(),
- ArrayRef<int64_t>{0});
- Value bcast = vector::BroadcastOp::create(rewriter, loc, valueTy, scalar);
- rewriter.replaceOp(op, bcast);
- return success();
-}
-
namespace {
/// Look up the subgroup size from the enclosing gpu.module's xevm.target.
@@ -1026,148 +985,133 @@ static unsigned lookupSubgroupSize(Operation *op) {
return uArch ? static_cast<unsigned>(uArch->getSubgroupSize()) : 16u;
}
-struct CoalesceLoadPattern final : OpRewritePattern<xegpu::LoadGatherOp> {
- CoalesceLoadPattern(MLIRContext *ctx, unsigned maxChunkSize,
- DataFlowSolver &solver)
- : OpRewritePattern(ctx), maxChunkSize(maxChunkSize), solver(solver) {}
-
- LogicalResult matchAndRewrite(xegpu::LoadGatherOp op,
- PatternRewriter &rewriter) const override {
- auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
- if (!offsetsTy)
- return rewriter.notifyMatchFailure(op, "expected vector offsets");
- if (offsetsTy.getNumElements() <= 1)
- return rewriter.notifyMatchFailure(op, "nothing to coalesce");
- auto valueTy = op.getValueType();
- if (!valueTy)
- return rewriter.notifyMatchFailure(op, "expected vector value");
- if (!isAllTrueMask(op.getMask()))
- return rewriter.notifyMatchFailure(op, "non-uniform mask");
-
- // Already coalesced (a previous run, or another pass tagged it).
- if (auto layout = op.getLayoutAttr())
- if (!layout.getEffectiveLaneDataAsInt().empty())
- return rewriter.notifyMatchFailure(op, "lane_data already set");
-
- int64_t origChunk = static_cast<int64_t>(op.getChunkSize().value_or(1));
- if (op.getChunkSizeAttr() && origChunk > 1)
- return rewriter.notifyMatchFailure(
- op, "explicit chunk_size > 1, leaving op alone");
-
- unsigned subgroupSize = lookupSubgroupSize(op);
-
- const auto *lat = solver.lookupState<AxisInfoLattice>(op.getOffsets());
- if (!lat || !lat->getValue().isInitialized())
- return rewriter.notifyMatchFailure(op, "no axis-info available");
-
- auto d = decide(lat->getValue(), offsetsTy.getShape(), origChunk,
- maxChunkSize, subgroupSize);
-
- // TODO: Don't do broadcast coalescing for now.
- // if (d.kind == CoalesceDecision::Kind::Broadcast)
- // return rewriteBroadcastLoad(op, rewriter);
-
- if (d.kind != CoalesceDecision::Kind::Chunked)
- return rewriter.notifyMatchFailure(op, "offsets not coalescible");
-
- auto layout = buildLaneDataLayout(op.getContext(), valueTy.getRank(),
- d.laneLayout, d.factor);
- // If the op carried an explicit chunk_size = 1 (the trivial / default
- // value), drop it: the new lane_data FCD > 1 supersedes it.
- bool dropChunk = op.getChunkSizeAttr() && origChunk == 1 && d.factor > 1;
- rewriter.modifyOpInPlace(op, [&] {
- op.setLayoutAttr(layout);
- if (dropChunk)
- op.removeChunkSizeAttr();
- });
- return success();
- }
+/// Discardable attribute name for the coalesce hint.
+static constexpr llvm::StringLiteral kCoalesceHintAttrName =
+ "xegpu.coalesce_hint";
- unsigned maxChunkSize;
- DataFlowSolver &solver;
-};
+/// Common analysis preconditions: vector offsets/value, all-true mask,
+/// no existing non-trivial lane_data, no explicit chunk_size > 1.
+template <typename OpTy>
+static bool isCandidateForCoalesce(OpTy op) {
+ auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
+ if (!offsetsTy || offsetsTy.getNumElements() <= 1)
+ return false;
+ if (!op.getValueType())
+ return false;
+ if (!isAllTrueMask(op.getMask()))
+ return false;
+ if (auto layout = op.getLayoutAttr())
+ if (!layout.getEffectiveLaneDataAsInt().empty())
+ return false;
+ if (op.getChunkSizeAttr() && op.getChunkSize().value_or(1) > 1)
+ return false;
+ return true;
+}
-struct CoalesceStorePattern final : OpRewritePattern<xegpu::StoreScatterOp> {
- CoalesceStorePattern(MLIRContext *ctx, unsigned maxChunkSize,
- DataFlowSolver &solver)
- : OpRewritePattern(ctx), maxChunkSize(maxChunkSize), solver(solver) {}
-
- LogicalResult matchAndRewrite(xegpu::StoreScatterOp op,
- PatternRewriter &rewriter) const override {
- auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
- if (!offsetsTy)
- return rewriter.notifyMatchFailure(op, "expected vector offsets");
- if (offsetsTy.getNumElements() <= 1)
- return rewriter.notifyMatchFailure(op, "nothing to coalesce");
- auto valueTy = op.getValueType();
- if (!valueTy)
- return rewriter.notifyMatchFailure(op, "expected vector value");
- if (!isAllTrueMask(op.getMask()))
- return rewriter.notifyMatchFailure(op, "non-uniform mask");
-
- if (auto layout = op.getLayoutAttr())
- if (!layout.getEffectiveLaneDataAsInt().empty())
- return rewriter.notifyMatchFailure(op, "lane_data already set");
-
- int64_t origChunk = static_cast<int64_t>(op.getChunkSize().value_or(1));
- if (op.getChunkSizeAttr() && origChunk > 1)
- return rewriter.notifyMatchFailure(
- op, "explicit chunk_size > 1, leaving op alone");
-
- unsigned subgroupSize = lookupSubgroupSize(op);
-
- const auto *lat = solver.lookupState<AxisInfoLattice>(op.getOffsets());
- if (!lat || !lat->getValue().isInitialized())
- return rewriter.notifyMatchFailure(op, "no axis-info available");
-
- auto d = decide(lat->getValue(), offsetsTy.getShape(), origChunk,
- maxChunkSize, subgroupSize);
-
- if (d.kind == CoalesceDecision::Kind::Broadcast)
- return rewriter.notifyMatchFailure(
- op, "all-equal offsets on store would be ambiguous");
-
- if (d.kind != CoalesceDecision::Kind::Chunked)
- return rewriter.notifyMatchFailure(op, "offsets not coalescible");
-
- auto layout = buildLaneDataLayout(op.getContext(), valueTy.getRank(),
- d.laneLayout, d.factor);
- bool dropChunk = op.getChunkSizeAttr() && origChunk == 1 && d.factor > 1;
- rewriter.modifyOpInPlace(op, [&] {
- op.setLayoutAttr(layout);
- if (dropChunk)
- op.removeChunkSizeAttr();
- });
- return success();
- }
+/// Run the analysis on a single op. If the offsets analyze as `Chunked`,
+/// stamp a `xegpu.coalesce_hint` attribute carrying the FCD lane_data
+/// factor.
+template <typename OpTy>
+static void analyzeAndStampHint(OpTy op, DataFlowSolver &solver,
+ unsigned maxChunkSize) {
+ if (!isCandidateForCoalesce(op))
+ return;
+ auto offsetsTy = cast<VectorType>(op.getOffsets().getType());
+ unsigned subgroupSize = lookupSubgroupSize(op);
+ const auto *lat = solver.lookupState<AxisInfoLattice>(op.getOffsets());
+ if (!lat || !lat->getValue().isInitialized())
+ return;
+ int64_t origChunk = static_cast<int64_t>(op.getChunkSize().value_or(1));
+ auto d = decide(lat->getValue(), offsetsTy.getShape(), origChunk,
+ maxChunkSize, subgroupSize);
+ if (d.kind != CoalesceDecision::Kind::Chunked)
+ return;
+ auto hint = xegpu::CoalesceHintAttr::get(op.getContext(), d.factor);
+ op->setAttr(kCoalesceHintAttrName, hint);
+}
- unsigned maxChunkSize;
- DataFlowSolver &solver;
-};
+/// Apply a stamped hint on `op`: build a lane_layout/lane_data/inst_data
+/// layout from the hint's `factor` and the op's offsets inner extent +
+/// chip-derived subgroup size, install it, drop a trivial `chunk_size = 1`,
+/// and remove the hint. Returns success on apply (or no-op when no hint),
+/// failure when the hint is malformed.
+template <typename OpTy>
+static LogicalResult applyHintOnOp(OpTy op) {
+ auto hint = op->template getAttrOfType<xegpu::CoalesceHintAttr>(
+ kCoalesceHintAttrName);
+ if (!hint)
+ return success(); // no hint: idempotent no-op.
-struct XeGPUCoalesceGatherScatterPass final
- : public xegpu::impl::XeGPUCoalesceGatherScatterBase<
- XeGPUCoalesceGatherScatterPass> {
- using XeGPUCoalesceGatherScatterBase::XeGPUCoalesceGatherScatterBase;
-
- void runOnOperation() override {
- Operation *root = getOperation();
-
- DataFlowSolver solver;
- solver.load<dataflow::DeadCodeAnalysis>();
- solver.load<mlir::xegpu::detail::axis_dataflow::AxisInfoAnalysis>();
- if (failed(solver.initializeAndRun(root)))
- return signalPassFailure();
-
- MLIRContext *ctx = &getContext();
- RewritePatternSet patterns(ctx);
- patterns.add<CoalesceLoadPattern, CoalesceStorePattern>(ctx, maxChunkSize,
- solver);
- if (failed(applyPatternsGreedily(root, std::move(patterns))))
- return signalPassFailure();
- }
-};
+ auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
+ auto valueTy = op.getValueType();
+ if (!offsetsTy || !valueTy || offsetsTy.getNumElements() <= 1)
+ return failure();
+
+ int64_t factor = hint.getFactor().getInt();
+ int64_t inner = offsetsTy.getShape().back();
+ unsigned subgroupSize = lookupSubgroupSize(op);
+ int64_t laneLayout =
+ largestPow2Divisor(inner, std::min<int64_t>(subgroupSize, inner));
+ if (laneLayout < 1 || inner % (laneLayout * factor) != 0)
+ return failure();
+
+ auto layout = buildLaneDataLayout(op.getContext(), valueTy.getRank(),
+ laneLayout, factor);
+ int64_t origChunk = static_cast<int64_t>(op.getChunkSize().value_or(1));
+ bool dropChunk = op.getChunkSizeAttr() && origChunk == 1 && factor > 1;
+ op.setLayoutAttr(layout);
+ if (dropChunk)
+ op.removeChunkSizeAttr();
+ op->removeAttr(kCoalesceHintAttrName);
+ return success();
+}
} // namespace
} // namespace
+
+//===----------------------------------------------------------------------===//
+// Public APIs.
+//===----------------------------------------------------------------------===//
+
+void mlir::xegpu::runCoalesceGatherScatterAnalysis(
+ Operation *root, const CoalesceGatherScatterAnalysisOptions &options) {
+ DataFlowSolver solver;
+ solver.load<dataflow::DeadCodeAnalysis>();
+ solver.load<mlir::xegpu::detail::axis_dataflow::AxisInfoAnalysis>();
+ if (failed(solver.initializeAndRun(root)))
+ return;
+
+ root->walk([&](Operation *op) {
+ if (auto load = dyn_cast<xegpu::LoadGatherOp>(op))
+ analyzeAndStampHint(load, solver, options.maxChunkSize);
+ else if (auto store = dyn_cast<xegpu::StoreScatterOp>(op))
+ analyzeAndStampHint(store, solver, options.maxChunkSize);
+ });
+}
+
+LogicalResult mlir::xegpu::applyCoalesceGatherScatterHint(Operation *op) {
+ if (auto load = dyn_cast<xegpu::LoadGatherOp>(op))
+ return applyHintOnOp(load);
+ if (auto store = dyn_cast<xegpu::StoreScatterOp>(op))
+ return applyHintOnOp(store);
+ // Hint attached to an unsupported op: silently drop it.
+ if (op->hasAttr(kCoalesceHintAttrName))
+ op->removeAttr(kCoalesceHintAttrName);
+ return success();
+}
+
+void mlir::xegpu::applyCoalesceGatherScatterHints(Operation *root) {
+ root->walk([&](Operation *op) {
+ if (op->hasAttr(kCoalesceHintAttrName))
+ (void)applyCoalesceGatherScatterHint(op);
+ });
+}
+
+void mlir::xegpu::clearCoalesceGatherScatterHints(Operation *root) {
+ StringRef name = kCoalesceHintAttrName;
+ root->walk([&](Operation *op) {
+ if (op->hasAttr(name))
+ op->removeAttr(name);
+ });
+}
diff --git a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
new file mode 100644
index 0000000000000..ddfcf6e506c94
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
@@ -0,0 +1,105 @@
+// RUN: mlir-opt -split-input-file \
+// RUN: -test-xegpu-coalesce-gather-scatter="analyze-only=true" %s | FileCheck %s
+
+// Analyze-only mode: stamps `xegpu.coalesce_hint` on coalescible ops and
+// leaves the layout / chunk_size unchanged. This test pins the hint
+// attribute contract that the apply API (and downstream propagator
+// integrations) consume.
+
+// -----
+// 1-D vector.step, fully coalescible -> hint with factor = 2 stamped.
+// CHECK-LABEL: func.func @load_step_offsets(
+// CHECK: xegpu.load
+// CHECK-SAME: <{chunk_size = 1 : i64}>
+// CHECK-SAME: {xegpu.coalesce_hint = #xegpu.coalesce_hint<factor = 2 : i64>}
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+// CHECK-NOT: lane_data
+func.func @load_step_offsets(%ptr: i64) -> vector<32xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// 2-D leading-1 dim: hint stamped on the load with factor = 2.
+// CHECK-LABEL: func.func @load_2d_leading_unit(
+// CHECK: xegpu.load
+// CHECK-SAME: {xegpu.coalesce_hint = #xegpu.coalesce_hint<factor = 2 : i64>}
+// CHECK-SAME: : i64, vector<1x32xindex>, vector<1x32xi1> -> vector<1x32xf32>
+// CHECK-NOT: lane_data
+func.func @load_2d_leading_unit(%ptr: i64) -> vector<1x32xf32> {
+ %step = vector.step : vector<32xindex>
+ %offsets = vector.shape_cast %step : vector<32xindex> to vector<1x32xindex>
+ %mask = arith.constant dense<true> : vector<1x32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<1x32xindex>, vector<1x32xi1> -> vector<1x32xf32>
+ return %v : vector<1x32xf32>
+}
+
+// -----
+// Stride-4 offsets: not coalescible, no hint stamped.
+// CHECK-LABEL: func.func @load_stride4_no_hint(
+// CHECK: xegpu.load
+// CHECK-NOT: coalesce_hint
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+func.func @load_stride4_no_hint(%ptr: i64) -> vector<32xf32> {
+ %c4 = arith.constant 4 : index
+ %step = vector.step : vector<32xindex>
+ %splat = vector.broadcast %c4 : index to vector<32xindex>
+ %offsets = arith.muli %step, %splat : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// All-equal offsets: classified as Broadcast by decide(); analysis returns
+// None at the stamping stage, so no hint is stamped (broadcast-load
+// rewrite was removed).
+// CHECK-LABEL: func.func @load_broadcast_offsets_no_hint(
+// CHECK: xegpu.load
+// CHECK-NOT: coalesce_hint
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+func.func @load_broadcast_offsets_no_hint(%ptr: i64) -> vector<32xf32> {
+ %offsets = arith.constant dense<0> : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// Explicit chunk_size > 1: precondition fails, no hint stamped (regardless
+// of analysis result).
+// CHECK-LABEL: gpu.func @load_explicit_chunk_no_hint(
+// CHECK: xegpu.load
+// CHECK-NOT: coalesce_hint
+// CHECK-SAME: <{chunk_size = 2 : i64}>
+gpu.module @kernel_explicit_chunk [#xevm.target<chip = "pvc">] {
+ gpu.func @load_explicit_chunk_no_hint(%ptr: i64) -> vector<32x2xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 2 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32x2xf32>
+ gpu.return %v : vector<32x2xf32>
+ }
+}
+
+// -----
+// Store with vector.step offsets: hint stamped with factor = 2 on the store
+// op.
+// CHECK-LABEL: func.func @store_step_hint(
+// CHECK: xegpu.store
+// CHECK-SAME: {xegpu.coalesce_hint = #xegpu.coalesce_hint<factor = 2 : i64>}
+// CHECK-SAME: : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
+// CHECK-NOT: lane_data
+func.func @store_step_hint(%ptr: i64, %v: vector<32xf32>) {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ xegpu.store %v, %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
+ return
+}
diff --git a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
index 859580c95906f..171879f18ad87 100644
--- a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
+++ b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
@@ -1,5 +1,5 @@
-// RUN: mlir-opt -split-input-file -xegpu-coalesce-gather-scatter %s | FileCheck %s
-// RUN: mlir-opt -split-input-file -xegpu-coalesce-gather-scatter="max-chunk-size=4" %s | FileCheck --check-prefix=CHECK4 %s
+// RUN: mlir-opt -split-input-file -test-xegpu-coalesce-gather-scatter %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -test-xegpu-coalesce-gather-scatter="max-chunk-size=4" %s | FileCheck --check-prefix=CHECK4 %s
// -----
// vector.step offsets -> stride 1, fully coalescible.
diff --git a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
index 7006051b9c033..048dde1e748cc 100644
--- a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
+++ b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
@@ -495,6 +495,48 @@ struct TestXeGPULayoutInterface
}
};
+struct TestXeGPUCoalesceGatherScatter
+ : public PassWrapper<TestXeGPUCoalesceGatherScatter, OperationPass<>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPUCoalesceGatherScatter)
+
+ StringRef getArgument() const final {
+ return "test-xegpu-coalesce-gather-scatter";
+ }
+
+ StringRef getDescription() const final {
+ return "Test the XeGPU coalesce-gather-scatter analysis + apply APIs.";
+ }
+
+ void getDependentDialects(::mlir::DialectRegistry ®istry) const override {
+ registry.insert<arith::ArithDialect>();
+ registry.insert<vector::VectorDialect>();
+ registry.insert<xegpu::XeGPUDialect>();
+ }
+
+ TestXeGPUCoalesceGatherScatter() = default;
+ TestXeGPUCoalesceGatherScatter(const TestXeGPUCoalesceGatherScatter &pass)
+ : PassWrapper(pass) {}
+
+ Option<unsigned> maxChunkSize{
+ *this, "max-chunk-size",
+ llvm::cl::desc("Upper bound on the produced lane_data FCD."),
+ llvm::cl::init(8)};
+
+ Option<bool> analyzeOnly{
+ *this, "analyze-only",
+ llvm::cl::desc("Only run the analysis (stamp xegpu.coalesce_hint "
+ "attributes); do not apply."),
+ llvm::cl::init(false)};
+
+ void runOnOperation() override {
+ xegpu::CoalesceGatherScatterAnalysisOptions options;
+ options.maxChunkSize = maxChunkSize;
+ xegpu::runCoalesceGatherScatterAnalysis(getOperation(), options);
+ if (!analyzeOnly)
+ xegpu::applyCoalesceGatherScatterHints(getOperation());
+ }
+};
+
} // namespace
namespace mlir {
@@ -509,6 +551,7 @@ void registerTestXeGPULowerings() {
PassRegistration<TestXeGPUPropagateLayouts>();
PassRegistration<TestXeGPUResolveLayoutConflicts>();
PassRegistration<TestXeGPUArrayLengthOptimization>();
+ PassRegistration<TestXeGPUCoalesceGatherScatter>();
}
} // namespace test
} // namespace mlir
More information about the Mlir-commits
mailing list