[Mlir-commits] [mlir] [mlir][XeGPU] Apply gather/scatter coalescing during lane layout propagation (PR #200979)
Md Abdullah Shahneous Bari
llvmlistbot at llvm.org
Mon Jun 1 17:53:56 PDT 2026
https://github.com/mshahneo created https://github.com/llvm/llvm-project/pull/200979
Integrates gather/scatter coalescing into xegpu-propagate-layout (lane level), so contiguous-per-lane `xegpu.load/xegpu.store` ops are distributed with a wider lane_data along the fastest-changing dim (FCD) — fewer, wider memory messages — without ever producing an unlowerable layout conversion.
Stacked on top of: 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/6] [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/6] [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/6] [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/6] [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/6] [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
>From 6e994df7121f066f008df7a31fdea3cfe97b5da1 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Tue, 2 Jun 2026 00:37:05 +0000
Subject: [PATCH 6/6] [mlir][XeGPU] Apply gather/scatter coalescing during lane
layout propagation.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Integrate the coalesce-gather-scatter analysis into
`xegpu-propagate-layout` at the lane level, replacing the standalone
pass that was removed from the pipeline in the previous commit.
`xegpu::propagateLayouts` now, when `layoutKind == Lane`:
1. Runs `runCoalesceGatherScatterAnalysis` up front to stamp
`xegpu.coalesce_hint` on coalescible xegpu.load/store ops.
2. Lets the hint flow through the existing backward layout
propagation (see below).
3. Strips any leftover hints on exit (scope guard).
The key insight is that propagation is a *backward* dataflow analysis,
so a post-hoc per-op conflict check cannot work: a load feeding a
reduction over the FCD and a load feeding an elementwise chain settle
to the same layout, and patching one op after the global fixpoint
reintroduces a lane_data mismatch with its already-fixed consumer.
Instead the hint is consumed *inside* anchor setup so the factor flows
through the existing `meet`/`min` machinery:
- Store visitor (`visitStoreScatterOp`): a coalescing store is a sink
free to choose its layout, so when it carries a hint its anchor
`lane_data[FCD]` is seeded to `factor`
(`setupStoreScatterAnchorLayout` / `setupGenericStoreAnchorLayout`
gain an optional `coalesceFactor`, default 1 = prior behavior).
- The seeded lane_data flows backward to producers via the operand
`meet`. The load visitor already caps its lane_data by the consumer
(`min(consumerLaneData, maxChunkSize)`), so a reduction consumer
(lane_data=1) overrides the hint and the load is not coalesced —
no unlowerable xegpu.convert_layout is created. An elementwise /
coalescing-store chain keeps the factor.
This fixes the IMEX `reduction_4D_3D_mean_intra_sg_intra_lane` kernel,
which previously failed to legalize a chunked-vs-round-robin
`xegpu.convert_layout` between a coalesced load and a reduction over
the inner dim; the kernel now lowers with zero convert_layout ops.
Also expose `getCoalesceHintAttrName()` from Transforms.h so the
propagator and the analysis share one attribute-name definition.
New lit test `propagate-layout-coalesce.mlir` covers the no-conflict
chain (load -> mulf -> store coalesces to lane_data=[2]), the reduction
barrier (load that would coalesce stays lane_data=[1,1], no convert),
and hint cleanup.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../Dialect/XeGPU/Transforms/Transforms.h | 5 ++
.../XeGPU/Transforms/XeGPULayoutImpl.h | 12 +--
.../Transforms/XeGPUCoalesceGatherScatter.cpp | 18 ++--
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 30 +++++--
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 32 ++++++-
.../XeGPU/propagate-layout-coalesce.mlir | 83 +++++++++++++++++++
6 files changed, 157 insertions(+), 23 deletions(-)
create mode 100644 mlir/test/Dialect/XeGPU/propagate-layout-coalesce.mlir
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
index eeaf0a8809ea7..8758d301fd34c 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
@@ -87,6 +87,11 @@ void populateXeGPUSgToWiDistributeTypeConversionAndLegality(
// Coalesce gather/scatter analysis + apply.
//===----------------------------------------------------------------------===//
+/// Discardable attribute name carrying the coalesce hint
+/// (`#xegpu.coalesce_hint<factor = N>`) stamped by
+/// `runCoalesceGatherScatterAnalysis`.
+inline StringRef getCoalesceHintAttrName() { return "xegpu.coalesce_hint"; }
+
/// Options controlling `runCoalesceGatherScatterAnalysis`.
struct CoalesceGatherScatterAnalysisOptions {
/// Upper bound on the per-lane chunk size produced by coalescing. Mirrors
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 8ffba7ff36208..8f1fd54bebdb6 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -205,11 +205,13 @@ setupLoadMatrixAnchorLayout(LayoutKind layoutKind, VectorType vectorTy,
DistributeLayoutAttr consumerLayout,
const uArch::uArch *uArch);
-/// Sets up the anchor layout for a store scatter operation.
-DistributeLayoutAttr setupStoreScatterAnchorLayout(LayoutKind layoutKind,
- VectorType vectorTy,
- int chunkSize,
- const uArch::uArch *uArch);
+/// Sets up the anchor layout for a store scatter operation. When
+/// `coalesceFactor > 1`, the store is a coalescing sink: its lane_data along
+/// the fastest-changing dim is seeded to `coalesceFactor` (and lane_layout
+/// shrunk to match), so the factor propagates backward to producers.
+DistributeLayoutAttr setupStoreScatterAnchorLayout(
+ LayoutKind layoutKind, VectorType vectorTy, int chunkSize,
+ const uArch::uArch *uArch, int coalesceFactor = 1);
/// Sets up the anchor layout for a store matrix operation.
DistributeLayoutAttr setupStoreMatrixAnchorLayout(LayoutKind layoutKind,
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
index c8ab290b9d654..44243f7292e67 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
@@ -985,10 +985,6 @@ static unsigned lookupSubgroupSize(Operation *op) {
return uArch ? static_cast<unsigned>(uArch->getSubgroupSize()) : 16u;
}
-/// Discardable attribute name for the coalesce hint.
-static constexpr llvm::StringLiteral kCoalesceHintAttrName =
- "xegpu.coalesce_hint";
-
/// Common analysis preconditions: vector offsets/value, all-true mask,
/// no existing non-trivial lane_data, no explicit chunk_size > 1.
template <typename OpTy>
@@ -1027,7 +1023,7 @@ static void analyzeAndStampHint(OpTy op, DataFlowSolver &solver,
if (d.kind != CoalesceDecision::Kind::Chunked)
return;
auto hint = xegpu::CoalesceHintAttr::get(op.getContext(), d.factor);
- op->setAttr(kCoalesceHintAttrName, hint);
+ op->setAttr(xegpu::getCoalesceHintAttrName(), hint);
}
/// Apply a stamped hint on `op`: build a lane_layout/lane_data/inst_data
@@ -1038,7 +1034,7 @@ static void analyzeAndStampHint(OpTy op, DataFlowSolver &solver,
template <typename OpTy>
static LogicalResult applyHintOnOp(OpTy op) {
auto hint = op->template getAttrOfType<xegpu::CoalesceHintAttr>(
- kCoalesceHintAttrName);
+ xegpu::getCoalesceHintAttrName());
if (!hint)
return success(); // no hint: idempotent no-op.
@@ -1062,7 +1058,7 @@ static LogicalResult applyHintOnOp(OpTy op) {
op.setLayoutAttr(layout);
if (dropChunk)
op.removeChunkSizeAttr();
- op->removeAttr(kCoalesceHintAttrName);
+ op->removeAttr(xegpu::getCoalesceHintAttrName());
return success();
}
@@ -1096,20 +1092,20 @@ LogicalResult mlir::xegpu::applyCoalesceGatherScatterHint(Operation *op) {
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);
+ if (op->hasAttr(xegpu::getCoalesceHintAttrName()))
+ op->removeAttr(xegpu::getCoalesceHintAttrName());
return success();
}
void mlir::xegpu::applyCoalesceGatherScatterHints(Operation *root) {
root->walk([&](Operation *op) {
- if (op->hasAttr(kCoalesceHintAttrName))
+ if (op->hasAttr(xegpu::getCoalesceHintAttrName()))
(void)applyCoalesceGatherScatterHint(op);
});
}
void mlir::xegpu::clearCoalesceGatherScatterHints(Operation *root) {
- StringRef name = kCoalesceHintAttrName;
+ StringRef name = xegpu::getCoalesceHintAttrName();
root->walk([&](Operation *op) {
if (op->hasAttr(name))
op->removeAttr(name);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 813e9608dbf8e..f4f6ff3d4e2ff 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1273,7 +1273,7 @@ static xegpu::DistributeLayoutAttr
setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
mlir::MLIRContext *context, bool isChunkedStore,
int maxChunkSize, ArrayRef<int64_t> srcShape,
- int subgroupSize) {
+ int subgroupSize, int coalesceFactor = 1) {
int srcShapeSize = srcShape.size();
SmallVector<int> instData(srcShapeSize, 1);
@@ -1287,13 +1287,27 @@ setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
}
if (!isChunkedStore) {
+ int64_t inner = srcShape.back();
if (layoutKind == xegpu::LayoutKind::InstData) {
instData[srcShapeSize - 1] =
- std::min(subgroupSize, static_cast<int>(srcShape.back()));
+ std::min(subgroupSize, static_cast<int>(inner));
return xegpu::LayoutAttr::get(context, instData);
} else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneLayout[srcShapeSize - 1] =
- std::min(subgroupSize, static_cast<int>(srcShape.back()));
+ // Coalescing sink: seed lane_data[FCD] = coalesceFactor (capped so
+ // lane_layout * lane_data divides the inner extent), then shrink
+ // lane_layout[FCD] accordingly. coalesceFactor == 1 reproduces the
+ // default (lane_layout = min(subgroupSize, inner), lane_data = 1).
+ int factor = std::max(1, coalesceFactor);
+ int laneLayoutInner = std::min<int64_t>(subgroupSize, inner / factor);
+ if (laneLayoutInner < 1)
+ laneLayoutInner = 1;
+ // Fall back to the trivial layout if the factor doesn't divide cleanly.
+ if (factor > 1 && inner % (laneLayoutInner * factor) != 0) {
+ factor = 1;
+ laneLayoutInner = std::min<int64_t>(subgroupSize, inner);
+ }
+ laneLayout[srcShapeSize - 1] = laneLayoutInner;
+ laneData[srcShapeSize - 1] = factor;
return xegpu::LayoutAttr::get(context, laneLayout, laneData);
}
} else {
@@ -1315,7 +1329,8 @@ setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
xegpu::DistributeLayoutAttr
xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
VectorType srcVecTy, int chunkSize,
- const uArch::uArch *uArch) {
+ const uArch::uArch *uArch,
+ int coalesceFactor) {
const int subgroupSize = uArch->getSubgroupSize();
ArrayRef<int64_t> srcShape = srcVecTy.getShape();
@@ -1326,8 +1341,11 @@ xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
int maxChunkSize = uArchInstruction->getMaxLaneStoreSize(elemBitWidth);
+ // A coalescing sink can carry at most `maxLaneStoreSize` per lane.
+ int factor = std::min(std::max(1, coalesceFactor), maxChunkSize);
return setupGenericStoreAnchorLayout(layoutKind, context, (chunkSize > 1),
- maxChunkSize, srcShape, subgroupSize);
+ maxChunkSize, srcShape, subgroupSize,
+ factor);
}
/// Sets up the anchor layout for a store matrix operation.
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index ffcda32e55c04..d6ce1d90f2956 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -15,6 +15,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/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
@@ -31,6 +32,7 @@
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/TypeSwitch.h"
@@ -1275,8 +1277,18 @@ void LayoutInfoPropagation::visitStoreScatterOp(
storeScatter.emitWarning("Not propagating, non-vector payload supplied.");
return;
}
+ // A coalesce hint (stamped by the coalesce-gather-scatter analysis at the
+ // start of lane propagation) seeds a non-trivial lane_data on the FCD so
+ // each lane stores `factor` contiguous elements. This factor flows
+ // backward to producers via the operand `meet` below.
+ int coalesceFactor = 1;
+ if (layoutKind == xegpu::LayoutKind::Lane) {
+ if (auto hint = storeScatter->getAttrOfType<xegpu::CoalesceHintAttr>(
+ xegpu::getCoalesceHintAttrName()))
+ coalesceFactor = static_cast<int>(hint.getFactor().getInt());
+ }
requiredAnchorLayoutAttr = xegpu::setupStoreScatterAnchorLayout(
- layoutKind, srcVecTy, chunkSize, uArch);
+ layoutKind, srcVecTy, chunkSize, uArch, coalesceFactor);
storeScatter.setLayoutAttr(requiredAnchorLayoutAttr);
}
@@ -1787,6 +1799,23 @@ struct XeGPUPropagateLayoutPass final
LogicalResult xegpu::propagateLayouts(OpBuilder &builder, Operation *target,
LayoutKind layoutKind,
unsigned indexBitWidth, bool printOnly) {
+ // At the lane level, run the coalesce-gather-scatter analysis up front so
+ // every coalescible xegpu.load/store carries an `xegpu.coalesce_hint`. The
+ // hint is consumed by the gather/scatter anchor-layout setup below: a
+ // coalescing store seeds a non-trivial lane_data on its FCD, which flows
+ // backward to producers via the operand `meet`. A consumer that requires a
+ // different lane_data (e.g. a reduction over the FCD) naturally overrides
+ // it, so no unlowerable convert_layout is created. Hints that propagation
+ // did not turn into a layout are stripped after the run.
+ bool coalesce = (layoutKind == LayoutKind::Lane) && !printOnly;
+ if (coalesce)
+ runCoalesceGatherScatterAnalysis(target);
+ auto cleanupFn = [&] {
+ if (coalesce)
+ clearCoalesceGatherScatterHints(target);
+ };
+ llvm::scope_exit<decltype(cleanupFn)> hintCleanup(std::move(cleanupFn));
+
RunLayoutInfoPropagation analysis(target, layoutKind, indexBitWidth);
// Print the analysis result and exit. (for debugging purposes)
if (printOnly) {
@@ -1876,6 +1905,7 @@ void XeGPUPropagateLayoutPass::runOnOperation() {
signalPassFailure();
return;
}
+
// Resolve layout conflicts if any.
if (failed(xegpu::resolveLayoutConflicts(getOperation()))) {
signalPassFailure();
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-coalesce.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-coalesce.mlir
new file mode 100644
index 0000000000000..c146cc45f9072
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-coalesce.mlir
@@ -0,0 +1,83 @@
+// RUN: mlir-opt -xevm-attach-target='chip=pvc' \
+// RUN: -test-xegpu-propagate-layouts="layout-kind=lane" -split-input-file %s \
+// RUN: | FileCheck %s
+
+// Lane-level layout propagation runs the coalesce-gather-scatter analysis up
+// front. A coalescing store seeds a non-trivial lane_data on its FCD, which
+// flows backward through the def-use chain to producers. A consumer that
+// needs a different lane_data (e.g. a reduction over the FCD) overrides the
+// hint, so no unlowerable xegpu.convert_layout is created.
+
+// -----
+// No-conflict path: load -> elementwise -> store, all contiguous on a
+// vector<32> (inner 32 > subgroup_size 16). The coalescing store seeds
+// lane_data = [2]; it flows back through the mulf to the load, so the whole
+// chain is laid out with lane_layout = [16], lane_data = [2].
+// CHECK-LABEL: gpu.func @coalesce_load_ew_store(
+// CHECK: xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>
+// CHECK: arith.mulf
+// CHECK-SAME: layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [2]>
+// CHECK: xegpu.store
+// CHECK-SAME: layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>
+// CHECK-NOT: xegpu.convert_layout
+gpu.module @kernel_chain {
+ gpu.func @coalesce_load_ew_store(%src: i64, %dst: i64) {
+ %step = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %src[%step], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ %c = arith.constant dense<2.0> : vector<32xf32>
+ %p = arith.mulf %v, %c : vector<32xf32>
+ xegpu.store %p, %dst[%step], %mask <{chunk_size = 1 : i64}>
+ : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
+ gpu.return
+ }
+}
+
+// -----
+// Conflict path: the load's offsets are contiguous on a vector<2x32> (inner
+// 32 > subgroup_size 16, so it would coalesce in isolation), but its only
+// consumer is a multi_reduction over the inner dim, which requires
+// lane_data = [1, 1] on the source. The reduction overrides the hint: the
+// load settles at lane_data = [1, 1], the reduction's source layout matches,
+// and no xegpu.convert_layout is created.
+// CHECK-LABEL: gpu.func @reduction_blocks_coalesce(
+// CHECK: xegpu.load
+// CHECK-SAME: layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+// CHECK: vector.multi_reduction
+// CHECK-SAME: #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>
+// CHECK-NOT: xegpu.convert_layout
+gpu.module @kernel_reduction {
+ gpu.func @reduction_blocks_coalesce(%src: i64, %dst: i64, %acc: vector<2xf32>) {
+ %step = vector.step : vector<64xindex>
+ %off2d = vector.shape_cast %step : vector<64xindex> to vector<2x32xindex>
+ %mask = arith.constant dense<true> : vector<2x32xi1>
+ %v = xegpu.load %src[%off2d], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<2x32xindex>, vector<2x32xi1> -> vector<2x32xf32>
+ %r = vector.multi_reduction <add>, %v, %acc [1]
+ : vector<2x32xf32> to vector<2xf32>
+ %roff = vector.step : vector<2xindex>
+ %rmask = arith.constant dense<true> : vector<2xi1>
+ xegpu.store %r, %dst[%roff], %rmask <{chunk_size = 1 : i64}>
+ : vector<2xf32>, i64, vector<2xindex>, vector<2xi1>
+ gpu.return
+ }
+}
+
+// -----
+// The coalesce hint is a transient analysis artifact: it must not survive
+// into the propagator's output.
+// CHECK-LABEL: gpu.func @no_leftover_hint(
+// CHECK-NOT: coalesce_hint
+gpu.module @kernel_cleanup {
+ gpu.func @no_leftover_hint(%src: i64, %dst: i64) {
+ %step = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %src[%step], %mask <{chunk_size = 1 : i64}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ xegpu.store %v, %dst[%step], %mask <{chunk_size = 1 : i64}>
+ : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
+ gpu.return
+ }
+}
More information about the Mlir-commits
mailing list