[Mlir-commits] [mlir] [mlir][XeGPU][Transform] Add gather/scatter coalescing analysis. (PR #201684)
Md Abdullah Shahneous Bari
llvmlistbot at llvm.org
Wed Jun 24 08:57:45 PDT 2026
https://github.com/mshahneo updated https://github.com/llvm/llvm-project/pull/201684
>From 1584887ae15b2dd831bf64630ed77a4a7f90dd68 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Thu, 4 Jun 2026 18:04:30 +0000
Subject: [PATCH 1/7] [mlir][XeGPU][Transform] Add gather/scatter coalescing
analysis.
Adds an AxisInfo-based dataflow analysis that classifies xegpu.load /
xegpu.store gather/scatter ops as coalescible and stamps a
`#xegpu.coalesce_hint<factor = N>` discardable attribute on them. The
analysis performs no rewrite of its own; the hint is consumed by a later
layout-propagation step (subsequent PRs) or by the apply helpers
(applyCoalesceGatherScatterHint) for testing.
Accesses tied to a vector.multi_reduction are not coalesced (isReductionTied):
consuming a coalesced lane_data[FCD] on a reduction needs reduction-aware
layout handling added in follow-up PRs; without it a coalesced store fed by a
reduction would force an unlowerable lane_data convert. Pure elementwise
gather/scatter is unaffected.
Includes:
- CoalesceHintAttr (XeGPUAttrs.td, XeGPUDialect.cpp).
- runCoalesceGatherScatterAnalysis / applyCoalesceGatherScatterHint(s) /
clearCoalesceGatherScatterHints public APIs (Transforms.h,
XeGPUCoalesceGatherScatter.cpp).
- A test-only pass (--test-xegpu-coalesce-gather-scatter, with
analyze-only) exercising the analysis and apply.
- Lit tests coalesce-gather-scatter{,-analyze}.mlir.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/IR/XeGPUAttrs.td | 41 +
.../Dialect/XeGPU/Transforms/Transforms.h | 44 +
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 15 +
.../Dialect/XeGPU/Transforms/CMakeLists.txt | 1 +
.../Transforms/XeGPUCoalesceGatherScatter.cpp | 1171 +++++++++++++++++
.../coalesce-gather-scatter-analyze.mlir | 105 ++
.../XeGPU/coalesce-gather-scatter.mlir | 476 +++++++
.../lib/Dialect/XeGPU/TestXeGPUTransforms.cpp | 43 +
8 files changed, 1896 insertions(+)
create mode 100644 mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
create mode 100644 mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
create mode 100644 mlir/test/Dialect/XeGPU/coalesce-gather-scatter.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/Transforms.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
index 919a69908bdce..6046f78dd59c0 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
@@ -80,6 +80,50 @@ void populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
TypeConverter &typeConverter, RewritePatternSet &patterns,
ConversionTarget &target);
+//===----------------------------------------------------------------------===//
+// 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
+ /// 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/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/CMakeLists.txt b/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
index 37922f7ef7d24..1b013400cff34 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
XeGPUSgToLaneDistribute.cpp
XeGPUUnroll.cpp
XeGPUWgToSgDistribute.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..669794f5c41fd
--- /dev/null
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
@@ -0,0 +1,1171 @@
+//===- 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
+// 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.
+//
+// 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.
+//
+//===----------------------------------------------------------------------===//
+
+#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"
+#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"
+#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/ADT/bit.h"
+#include "llvm/Support/MathExtras.h"
+#include <numeric>
+#include <optional>
+
+#define DEBUG_TYPE "xegpu-coalesce-gather-scatter"
+
+using namespace mlir;
+
+// 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 {
+
+//===----------------------------------------------------------------------===//
+// AxisInfo: per-axis contiguity / constancy / divisibility lattice.
+//===----------------------------------------------------------------------===//
+
+/// 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.
+/// - `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,
+/// 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;
+
+ 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 && innerStride == rhs.innerStride;
+ }
+
+ /// 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;
+ if (lhs.innerStride && rhs.innerStride &&
+ *lhs.innerStride == *rhs.innerStride)
+ out.innerStride = lhs.innerStride;
+ 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;
+ if (innerStride)
+ os << " innerStride=" << *innerStride;
+ }
+};
+
+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;
+ v.innerStride = 0;
+ 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 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))
+ 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};
+ v.innerStride = 1;
+ 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()) {
+ 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();
+ }
+ 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 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);
+ }
+ }
+ 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;
+ if (innerStride != std::numeric_limits<int64_t>::min())
+ v.innerStride = innerStride;
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+
+ // 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;
+ // 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();
+ }
+
+ // 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();
+ }
+ AxisInfo src = operands[0]->getValue();
+ if (!src.isInitialized()) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ // 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];
+ }
+ } 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;
+ // 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();
+ }
+
+ 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]);
+ }
+ // 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();
+ }
+
+ 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;
+ }
+ // 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();
+ }
+
+ 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, 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
+};
+
+/// 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.
+///
+/// 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,
+ unsigned subgroupSize) {
+ CoalesceDecision d;
+ if (!info.isInitialized() || offsetsShape.empty())
+ return d;
+ unsigned innerDim = offsetsShape.size() - 1;
+ int64_t inner = offsetsShape[innerDim];
+ if (inner < 2)
+ return d;
+
+ // 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.
+ 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(perLane, bound);
+ if (factor < 2)
+ return d;
+ d.kind = CoalesceDecision::Kind::Chunked;
+ d.laneLayout = laneLayout;
+ d.factor = factor;
+ return d;
+}
+
+/// 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();
+}
+
+/// 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 innerLaneLayout,
+ int64_t innerLaneData) {
+ SmallVector<int32_t> laneLayout(rank, 1);
+ SmallVector<int32_t> laneData(rank, 1);
+ 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.
+//===----------------------------------------------------------------------===//
+
+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;
+}
+
+/// 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;
+}
+
+/// True when `op` (a gather load or scatter store) is tied to a
+/// `vector.multi_reduction`: a load whose result feeds a reduction, or a
+/// store whose stored value comes from one (through layout-neutral /
+/// elementwise / insert glue).
+///
+/// Coalescing such an access is gated off for now: the analysis only sets
+/// `lane_data[FCD]`, and consuming that on a reduction requires reduction-
+/// aware layout handling that is not part of this change. A coalesced store
+/// fed (transitively) by a reduction would seed `lane_data[FCD] = N` that the
+/// reduction result (kept at `lane_data = 1`) cannot match, producing an
+/// unlowerable `xegpu.convert_layout lane_data=[1]<->[N]`. Follow-up PRs that
+/// add reduction coalescing relax this gate.
+template <typename OpTy>
+static bool isReductionTied(OpTy op) {
+ if constexpr (std::is_same_v<OpTy, xegpu::StoreScatterOp>) {
+ // Backward walk of the stored value's slice through the reassembly
+ // (convert_layout / shape_cast / insert(_strided_slice) / elementwise)
+ // that typically sits between a reduction and the store.
+ SmallVector<Value, 8> worklist{op.getValue()};
+ llvm::SmallPtrSet<Operation *, 16> seen;
+ unsigned steps = 0;
+ while (!worklist.empty() && steps++ < 64) {
+ Value v = worklist.pop_back_val();
+ Operation *def = v.getDefiningOp();
+ if (!def || !seen.insert(def).second)
+ continue;
+ if (isa<vector::MultiDimReductionOp>(def))
+ return true;
+ if (isa<vector::ShapeCastOp, vector::BitCastOp, xegpu::ConvertLayoutOp,
+ vector::InsertOp, vector::InsertStridedSliceOp>(def) ||
+ OpTrait::hasElementwiseMappableTraits(def))
+ for (Value operand : def->getOperands())
+ if (isa<VectorType>(operand.getType()))
+ worklist.push_back(operand);
+ }
+ return false;
+ } else {
+ for (Operation *user : op->getUsers()) {
+ Operation *u = user;
+ while (
+ u &&
+ isa<vector::ShapeCastOp, vector::BitCastOp, xegpu::ConvertLayoutOp>(
+ u)) {
+ if (u->getNumResults() != 1 || u->getResult(0).use_empty())
+ break;
+ u = *u->getResult(0).getUsers().begin();
+ }
+ if (u && isa<vector::MultiDimReductionOp>(u))
+ return true;
+ }
+ return false;
+ }
+}
+
+/// 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;
+ // Do not coalesce accesses tied to a reduction (see isReductionTied);
+ // reduction coalescing is added in follow-up PRs.
+ if (isReductionTied(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(xegpu::getCoalesceHintAttrName(), hint);
+}
+
+/// 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>(
+ xegpu::getCoalesceHintAttrName());
+ if (!hint)
+ return success(); // no hint: idempotent no-op.
+
+ 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(xegpu::getCoalesceHintAttrName());
+ 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(xegpu::getCoalesceHintAttrName()))
+ op->removeAttr(xegpu::getCoalesceHintAttrName());
+ return success();
+}
+
+void mlir::xegpu::applyCoalesceGatherScatterHints(Operation *root) {
+ root->walk([&](Operation *op) {
+ if (op->hasAttr(xegpu::getCoalesceHintAttrName()))
+ (void)applyCoalesceGatherScatterHint(op);
+ });
+}
+
+void mlir::xegpu::clearCoalesceGatherScatterHints(Operation *root) {
+ StringRef name = xegpu::getCoalesceHintAttrName();
+ 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
new file mode 100644
index 0000000000000..171879f18ad87
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
@@ -0,0 +1,476 @@
+// 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.
+// 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-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> {
+ %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: 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<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>
+ %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: func.func @load_dense_ap_offsets(
+// 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<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>
+}
+
+// -----
+// 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
+ %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: 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>
+ %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: func.func @store_step_offsets(
+// CHECK: xegpu.store
+// 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>
+ %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: 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>
+ %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: func.func @load_memref_step(
+// CHECK: xegpu.load
+// 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>
+ %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>
+}
+
+// -----
+// 2-D offsets with leading unit dim: inner dim treated as lane dim.
+// 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<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>
+ %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, 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-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<[
+ [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>
+}
+
+// -----
+// 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<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>
+ }
+}
+
+// -----
+// 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>
+}
+
+// -----
+// "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 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: 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<[
+ 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: same broadcast case
+// as remui above, currently disabled, so the load is left alone.
+// CHECK-LABEL: gpu.func @load_andi_inner_uniform(
+// 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<[
+ 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<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>
+ %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
+}
+
+// -----
+// 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>
+ }
+}
diff --git a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
index 5c3721630837d..b71bd6a0be618 100644
--- a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
+++ b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
@@ -446,6 +446,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 {
@@ -458,6 +500,7 @@ void registerTestXeGPULowerings() {
PassRegistration<TestXeGPUPropagateLayouts>();
PassRegistration<TestXeGPUResolveLayoutConflicts>();
PassRegistration<TestXeGPUArrayLengthOptimization>();
+ PassRegistration<TestXeGPUCoalesceGatherScatter>();
}
} // namespace test
} // namespace mlir
>From 5ef04fb11afa1215b7f1da6f99099863da88e764 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Wed, 10 Jun 2026 20:05:21 +0000
Subject: [PATCH 2/7] Address review comments.
---
.../mlir/Dialect/XeGPU/IR/XeGPUAttrs.td | 26 +++--
.../Dialect/XeGPU/Transforms/Transforms.h | 29 +++---
.../Transforms/XeGPUCoalesceGatherScatter.cpp | 99 ++++++++++---------
.../coalesce-gather-scatter-analyze.mlir | 33 ++-----
.../XeGPU/coalesce-gather-scatter.mlir | 74 +++++---------
.../lib/Dialect/XeGPU/TestXeGPUTransforms.cpp | 2 +-
6 files changed, 107 insertions(+), 156 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
index 50b67cefafa31..292302a74934c 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
@@ -933,24 +933,22 @@ def XeGPU_CoalesceHintAttr : XeGPUAttr<"CoalesceHint", "coalesce_hint"> {
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.
+ whose offsets describe a coalescible access pattern. `factor` is the
+ number of contiguous elements grouped together along the innermost
+ (fastest-changing) dimension of the access.
+
+ The attribute is purely advisory: it records the analysis's decision
+ without committing to a layout. A consumer turns `factor` into a concrete
+ `xegpu.layout` (deriving the `lane_layout` / `lane_data` split from the
+ op's offsets inner extent and the chip's subgroup size) and removes the
+ hint. A consumer that decides not to act on the hint (for example, a
+ propagator that detects a conflict with an anchor-driven layout) should
+ likewise 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>}>
+ {xegpu.coalesce_hint = #xegpu.coalesce_hint<factor = 4>}
: i64, vector<64xindex>, vector<64xi1> -> vector<64xf32>
```
}];
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
index 6046f78dd59c0..fb560acc3dc48 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
@@ -91,33 +91,26 @@ inline StringRef getCoalesceHintAttrName() { return "xegpu.coalesce_hint"; }
/// 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.
+ /// Upper bound on the number of contiguous elements grouped per lane 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).
+/// `lane_data` or a non-uniform mask are skipped (no hint stamped). The hint
+/// is consumed downstream — either by `coalesceGatherScatter` or by a
+/// layout-propagation pass that reads the hint directly.
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 turn every stamped `xegpu.coalesce_hint` into an
+/// equivalent `lane_layout` / `lane_data` / `inst_data` layout on its op,
+/// then remove the hint. A hint on an op that cannot be coalesced (e.g. one
+/// that isn't a gather/scatter) is simply dropped. Pairs with
+/// `runCoalesceGatherScatterAnalysis`.
+void coalesceGatherScatter(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
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
index 669794f5c41fd..52f1b321be900 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
@@ -6,14 +6,18 @@
//
//===----------------------------------------------------------------------===//
//
-// This pass coalesces neighbouring lanes of `xegpu.load` / `xegpu.store` 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`.
+// This file implements the gather/scatter coalescing analysis. It decides
+// whether an `xegpu.load` / `xegpu.store` gather/scatter accesses `N`
+// contiguous elements per lane along the innermost dimension, and, if so,
+// stamps a `xegpu.coalesce_hint<factor = N>` attribute recording the chosen
+// factor. The decision is driven by a small XeGPU-local axis-info dataflow
+// analysis tracking per-axis `contiguity`, `constancy`, and `divisibility`.
+//
+// The analysis performs no rewrite. The hint it stamps is turned into a
+// `lane_data` layout by `coalesceGatherScatter` (used by the test pass) or
+// read directly by layout propagation; 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
@@ -80,10 +84,15 @@ static constexpr int64_t kAxisInfoTop = 1LL << 30;
/// - `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.
+/// is an arithmetic progression with this (constant) step. This is the
+/// general form of which `contiguity`/`constancy` are the two special
+/// cases: `innerStride = 1` is the stride-1 contiguous case (and implies
+/// `contiguity[innerDim] > 1`), `innerStride = 0` is the all-equal case
+/// (implies `constancy[innerDim] > 1`), and any other value (e.g. 4) is a
+/// strided progression that neither `contiguity` nor `constancy` can
+/// represent — both read 1 there. Per-row base may differ across outer
+/// indices; per-row alignment is captured by `divisibility[innerDim]`.
+/// The coalescing decision only acts on the `innerStride = 1` case.
///
/// Pessimistic / entry value: contiguity=1, constancy=1, divisibility=1,
/// innerStride absent.
@@ -888,11 +897,10 @@ static int64_t largestPow2Divisor(int64_t numLanes, int64_t bound) {
/// `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.
+/// `maxChunkSize` and the offsets contiguity.
static CoalesceDecision decide(const AxisInfo &info,
ArrayRef<int64_t> offsetsShape,
- int64_t origChunk, unsigned maxChunkSize,
- unsigned subgroupSize) {
+ unsigned maxChunkSize, unsigned subgroupSize) {
CoalesceDecision d;
if (!info.isInitialized() || offsetsShape.empty())
return d;
@@ -918,9 +926,7 @@ static CoalesceDecision decide(const AxisInfo &info,
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;
+ int64_t budget = static_cast<int64_t>(maxChunkSize);
if (budget < 2)
return d;
@@ -969,11 +975,9 @@ static xegpu::LayoutAttr buildLaneDataLayout(MLIRContext *ctx, unsigned rank,
}
//===----------------------------------------------------------------------===//
-// Rewrites.
+// Analysis driver + hint apply.
//===----------------------------------------------------------------------===//
-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
@@ -986,7 +990,7 @@ static unsigned lookupSubgroupSize(Operation *op) {
}
/// Common analysis preconditions: vector offsets/value, all-true mask,
-/// no existing non-trivial lane_data, no explicit chunk_size > 1.
+/// no existing non-trivial lane_data.
template <typename OpTy>
static bool isCandidateForCoalesce(OpTy op) {
auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
@@ -999,8 +1003,6 @@ static bool isCandidateForCoalesce(OpTy op) {
if (auto layout = op.getLayoutAttr())
if (!layout.getEffectiveLaneDataAsInt().empty())
return false;
- if (op.getChunkSizeAttr() && op.getChunkSize().value_or(1) > 1)
- return false;
return true;
}
@@ -1075,9 +1077,8 @@ static void analyzeAndStampHint(OpTy op, DataFlowSolver &solver,
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);
+ auto d = decide(lat->getValue(), offsetsTy.getShape(), maxChunkSize,
+ subgroupSize);
if (d.kind != CoalesceDecision::Kind::Chunked)
return;
auto hint = xegpu::CoalesceHintAttr::get(op.getContext(), d.factor);
@@ -1086,9 +1087,9 @@ static void analyzeAndStampHint(OpTy op, 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.
+/// chip-derived subgroup size, install it, 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>(
@@ -1111,16 +1112,25 @@ static LogicalResult applyHintOnOp(OpTy op) {
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(xegpu::getCoalesceHintAttrName());
return success();
}
-} // namespace
+/// Apply a stamped `xegpu.coalesce_hint` on a single op. A hint on an op
+/// that cannot be coalesced (anything other than a gather/scatter) is just
+/// dropped. Returns failure only when the hint is malformed for a real
+/// gather/scatter.
+static LogicalResult applyHintOnOp(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(xegpu::getCoalesceHintAttrName()))
+ op->removeAttr(xegpu::getCoalesceHintAttrName());
+ return success();
+}
} // namespace
@@ -1136,6 +1146,10 @@ void mlir::xegpu::runCoalesceGatherScatterAnalysis(
if (failed(solver.initializeAndRun(root)))
return;
+ // The solver computed AxisInfo for the whole region in the single
+ // `initializeAndRun` above; offsets shared by several gather/scatter ops are
+ // analyzed only once. This walk is just per-op point lookups into that
+ // result (no re-analysis), turning each cached fact into a hint.
root->walk([&](Operation *op) {
if (auto load = dyn_cast<xegpu::LoadGatherOp>(op))
analyzeAndStampHint(load, solver, options.maxChunkSize);
@@ -1144,21 +1158,10 @@ void mlir::xegpu::runCoalesceGatherScatterAnalysis(
});
}
-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(xegpu::getCoalesceHintAttrName()))
- op->removeAttr(xegpu::getCoalesceHintAttrName());
- return success();
-}
-
-void mlir::xegpu::applyCoalesceGatherScatterHints(Operation *root) {
+void mlir::xegpu::coalesceGatherScatter(Operation *root) {
root->walk([&](Operation *op) {
if (op->hasAttr(xegpu::getCoalesceHintAttrName()))
- (void)applyCoalesceGatherScatterHint(op);
+ (void)applyHintOnOp(op);
});
}
diff --git a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
index ddfcf6e506c94..4b528704607a9 100644
--- a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
+++ b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
@@ -2,22 +2,20 @@
// 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.
+// leaves the layout 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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
return %v : vector<32xf32>
}
@@ -33,7 +31,7 @@ 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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<1x32xindex>, vector<1x32xi1> -> vector<1x32xf32>
return %v : vector<1x32xf32>
}
@@ -50,7 +48,7 @@ func.func @load_stride4_no_hint(%ptr: i64) -> vector<32xf32> {
%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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
return %v : vector<32xf32>
}
@@ -66,28 +64,11 @@ func.func @load_stride4_no_hint(%ptr: i64) -> 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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: 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.
@@ -99,7 +80,7 @@ gpu.module @kernel_explicit_chunk [#xevm.target<chip = "pvc">] {
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}>
+ xegpu.store %v, %ptr[%offsets], %mask
: 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 171879f18ad87..25237df7f8014 100644
--- a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
+++ b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter.mlir
@@ -5,19 +5,16 @@
// vector.step offsets -> stride 1, fully coalescible.
// 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-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> {
%offsets = vector.step : vector<32xindex>
%mask = arith.constant dense<true> : vector<32xi1>
- %v = xegpu.load %ptr[%offsets], %mask <{chunk_size = 1 : i64}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
return %v : vector<32xf32>
}
@@ -33,7 +30,7 @@ func.func @load_step_offsets(%ptr: i64) -> 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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
return %v : vector<32xf32>
}
@@ -50,7 +47,7 @@ func.func @load_dense_ap_offsets(%ptr: i64) -> vector<32xi32> {
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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<32xindex>, vector<32xi1> -> vector<32xi32>
return %v : vector<32xi32>
}
@@ -67,7 +64,7 @@ func.func @load_stride4_unchanged(%ptr: i64) -> vector<32xf32> {
%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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
return %v : vector<32xf32>
}
@@ -80,7 +77,7 @@ func.func @load_stride4_unchanged(%ptr: i64) -> vector<32xf32> {
// 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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
return %v : vector<32xf32>
}
@@ -94,7 +91,7 @@ func.func @load_partial_mask_unchanged(%ptr: i64, %mask: vector<32xi1>) -> vecto
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}>
+ xegpu.store %v, %ptr[%offsets], %mask
: vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
return
}
@@ -108,7 +105,7 @@ func.func @store_step_offsets(%ptr: i64, %v: vector<32xf32>) {
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}>
+ xegpu.store %v, %ptr[%offsets], %mask
: vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
return
}
@@ -122,7 +119,7 @@ func.func @store_broadcast_offsets_unchanged(%ptr: i64, %v: 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}>
+ %v = xegpu.load %m[%offsets], %mask
: memref<1024xf32>, vector<32xindex>, vector<32xi1> -> vector<32xf32>
return %v : vector<32xf32>
}
@@ -139,7 +136,7 @@ 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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<1x32xindex>, vector<1x32xi1> -> vector<1x32xf32>
return %v : vector<1x32xf32>
}
@@ -159,7 +156,7 @@ func.func @load_2d_dense_ap(%ptr: i64) -> vector<2x16xf32> {
[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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<2x16xindex>, vector<2x16xi1> -> vector<2x16xf32>
return %v : vector<2x16xf32>
}
@@ -168,7 +165,7 @@ func.func @load_2d_dense_ap(%ptr: i64) -> vector<2x16xf32> {
// 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.
+// budget = max-chunk-size = 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<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>
@@ -182,7 +179,7 @@ gpu.module @kernel [#xevm.target<chip = "pvc">] {
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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<2x32xindex>, vector<2x32xi1> -> vector<2x32xf32>
gpu.return %v : vector<2x32xf32>
}
@@ -202,7 +199,7 @@ gpu.module @kernel_2x32 [#xevm.target<chip = "pvc">] {
%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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<2x32xindex>, vector<2x32xi1> -> vector<2x32xf32>
gpu.return %v : vector<2x32xf32>
}
@@ -221,7 +218,7 @@ func.func @load_2d_non_ap_unchanged(%ptr: i64) -> vector<2x16xf32> {
[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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<2x16xindex>, vector<2x16xi1> -> vector<2x16xf32>
return %v : vector<2x16xf32>
}
@@ -255,7 +252,7 @@ gpu.module @kernel_reduction [#xevm.target<chip = "pvc">] {
: 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}>
+ %v = xegpu.load %ptr[%off], %mask
: i64, vector<2x32xindex>, vector<2x32xi1> -> vector<2x32xf32>
gpu.return %v : vector<2x32xf32>
}
@@ -279,7 +276,7 @@ gpu.module @kernel_divui [#xevm.target<chip = "pvc">] {
%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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
gpu.return %v : vector<32xf32>
}
@@ -301,7 +298,7 @@ gpu.module @kernel_divsi [#xevm.target<chip = "pvc">] {
%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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
gpu.return %v : vector<32xf32>
}
@@ -311,7 +308,7 @@ 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 decision picks
// the broadcast case, which is currently disabled, so the load is left
-// alone (no layout, no chunk_size change).
+// alone (no layout change).
// CHECK-LABEL: gpu.func @load_remui_inner_uniform(
// CHECK: xegpu.load
// CHECK-NOT: lane_data
@@ -325,7 +322,7 @@ gpu.module @kernel_remui [#xevm.target<chip = "pvc">] {
// (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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
gpu.return %v : vector<16xf32>
}
@@ -347,7 +344,7 @@ gpu.module @kernel_andi [#xevm.target<chip = "pvc">] {
%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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
gpu.return %v : vector<16xf32>
}
@@ -367,7 +364,7 @@ gpu.module @kernel_shli [#xevm.target<chip = "pvc">] {
%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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
gpu.return %v : vector<32xf32>
}
@@ -388,7 +385,7 @@ gpu.module @kernel_shli_shrui [#xevm.target<chip = "pvc">] {
%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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
gpu.return %v : vector<32xf32>
}
@@ -412,7 +409,7 @@ gpu.module @kernel_select [#xevm.target<chip = "pvc">] {
%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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
gpu.return %v : vector<32xf32>
}
@@ -433,7 +430,7 @@ gpu.module @kernel_divui_neg [#xevm.target<chip = "pvc">] {
%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}>
+ %v = xegpu.load %ptr[%offsets], %mask
: i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
gpu.return %v : vector<16xf32>
}
@@ -449,28 +446,7 @@ 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}>
+ xegpu.store %v, %ptr[%offsets], %mask
: 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>
- }
-}
diff --git a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
index b71bd6a0be618..f24b4dccaab04 100644
--- a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
+++ b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
@@ -484,7 +484,7 @@ struct TestXeGPUCoalesceGatherScatter
options.maxChunkSize = maxChunkSize;
xegpu::runCoalesceGatherScatterAnalysis(getOperation(), options);
if (!analyzeOnly)
- xegpu::applyCoalesceGatherScatterHints(getOperation());
+ xegpu::coalesceGatherScatter(getOperation());
}
};
>From fbaaaf14c113fbfa0b18f7907fd87c3c5483cea0 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Wed, 10 Jun 2026 20:25:51 +0000
Subject: [PATCH 3/7] Fix clang-format issue.
---
.../Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
index 52f1b321be900..788a5774afc91 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
@@ -1077,8 +1077,8 @@ static void analyzeAndStampHint(OpTy op, DataFlowSolver &solver,
const auto *lat = solver.lookupState<AxisInfoLattice>(op.getOffsets());
if (!lat || !lat->getValue().isInitialized())
return;
- auto d = decide(lat->getValue(), offsetsTy.getShape(), maxChunkSize,
- subgroupSize);
+ auto d =
+ decide(lat->getValue(), offsetsTy.getShape(), maxChunkSize, subgroupSize);
if (d.kind != CoalesceDecision::Kind::Chunked)
return;
auto hint = xegpu::CoalesceHintAttr::get(op.getContext(), d.factor);
>From 46cdcd81726941549fdc5ce87b29a1a76fa23576 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Thu, 18 Jun 2026 15:47:20 +0000
Subject: [PATCH 4/7] [mlir][XeGPU] Make coalesce_hint an optional op attribute
(user-authorable).
Promotes `coalesce_hint` from a loose discardable attribute to an inherent
optional attribute of `xegpu.load` / `xegpu.store` (`OptionalAttr<
XeGPU_CoalesceHintAttr>`). This gives a second source for the hint: a user
(or a higher-level frontend that already knows the access is
contiguous-per-lane) can set it directly, in which case the
coalesce-gather-scatter analysis is bypassed for that op.
- The analysis now yields to a pre-existing hint: analyzeAndStampHint skips
an op that already carries `coalesce_hint` (user-set, or stamped by an
earlier propagation level). This also makes it idempotent across the lane
and inst-data runs. It uses the generated typed accessors
(get/set/removeCoalesceHintAttr) instead of string-keyed discardable-attr
access; getCoalesceHintAttrName() is removed.
- Verification: the attribute verifier keeps `factor` a power of two >= 2;
a new op-level check (isValidCoalesceHint) requires the innermost offsets
dimension to be a multiple of `factor`. The `lane_layout * factor` split
is chip-dependent and is still checked at apply time.
- In-tree positional builders/creators of these ops pass the new trailing
optional attribute explicitly.
Tests: ops.mlir round-trips a user-set coalesce_hint on load and store;
invalid.mlir covers the three rejections (non-power-of-two, < 2, non-dividing
factor); coalesce-gather-scatter-analyze.mlir adds a user-hint-preserved case.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/IR/XeGPUAttrs.td | 36 ++++++++++--------
.../include/mlir/Dialect/XeGPU/IR/XeGPUOps.td | 6 ++-
.../Dialect/XeGPU/Transforms/Transforms.h | 5 ---
.../VectorToXeGPU/VectorToXeGPU.cpp | 8 ++--
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp | 37 +++++++++++++++++--
.../Transforms/XeGPUCoalesceGatherScatter.cpp | 33 +++++++++--------
.../Transforms/XeGPUSgToLaneDistribute.cpp | 5 ++-
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 6 ++-
.../Transforms/XeGPUWgToSgDistribute.cpp | 5 ++-
.../coalesce-gather-scatter-analyze.mlir | 22 +++++++++--
mlir/test/Dialect/XeGPU/invalid.mlir | 24 ++++++++++++
mlir/test/Dialect/XeGPU/ops.mlir | 17 +++++++++
12 files changed, 148 insertions(+), 56 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
index 292302a74934c..e16ee5031785c 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
@@ -928,27 +928,31 @@ 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 summary = [{Optional request to coalesce a gather/scatter access.}];
let description = [{
- `CoalesceHintAttr` is a discardable attribute attached by the
- coalesce-gather-scatter analysis to `xegpu.load` / `xegpu.store` ops
- whose offsets describe a coalescible access pattern. `factor` is the
- number of contiguous elements grouped together along the innermost
- (fastest-changing) dimension of the access.
-
- The attribute is purely advisory: it records the analysis's decision
- without committing to a layout. A consumer turns `factor` into a concrete
- `xegpu.layout` (deriving the `lane_layout` / `lane_data` split from the
- op's offsets inner extent and the chip's subgroup size) and removes the
- hint. A consumer that decides not to act on the hint (for example, a
- propagator that detects a conflict with an anchor-driven layout) should
- likewise remove the attribute rather than leave it dangling.
+ `CoalesceHintAttr` is the optional `coalesce_hint` attribute of
+ `xegpu.load` / `xegpu.store`. It requests that the access group `factor`
+ contiguous elements per lane along the innermost (fastest-changing)
+ dimension of the offsets. `factor` must be a power of two `>= 2`, and the
+ innermost offsets dimension must be a multiple of it.
+
+ The hint has two sources:
+ - it can be set directly by the user (or a higher-level frontend that
+ already knows the access is contiguous-per-lane), in which case the
+ coalesce-gather-scatter analysis is not consulted for that op; or
+ - it is stamped by the analysis when it proves the access coalescible.
+
+ The hint is a *request*, not a commitment: a consumer turns `factor` into
+ a concrete `xegpu.layout` (deriving the `lane_layout` / `lane_data` split
+ from the op's offsets inner extent and the chip's subgroup size) and
+ clears the hint. A consumer that declines (for example, layout propagation
+ detecting a conflict with an anchor-driven layout) likewise clears it
+ rather than leaving it dangling.
Example:
```mlir
- %v = xegpu.load %ptr[%offsets], %mask
- {xegpu.coalesce_hint = #xegpu.coalesce_hint<factor = 4>}
+ %v = xegpu.load %ptr[%offsets], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 4>}>
: i64, vector<64xindex>, vector<64xi1> -> vector<64xf32>
```
}];
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index af8c30742bcb2..a58ee7430a56d 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -773,7 +773,8 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
OptionalAttr<XeGPU_CacheHintAttr>:$l1_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l3_hint,
- OptionalAttr<DistributeLayoutAttr>:$layout);
+ OptionalAttr<DistributeLayoutAttr>:$layout,
+ OptionalAttr<XeGPU_CoalesceHintAttr>:$coalesce_hint);
let results = (outs XeGPU_ValueOrScalarType:$value);
let extraClassDeclaration = extraBaseClassDeclaration # [{
@@ -903,7 +904,8 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
OptionalAttr<XeGPU_CacheHintAttr>:$l1_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l3_hint,
- OptionalAttr<DistributeLayoutAttr>:$layout);
+ OptionalAttr<DistributeLayoutAttr>:$layout,
+ OptionalAttr<XeGPU_CoalesceHintAttr>:$coalesce_hint);
let extraClassDeclaration = extraBaseClassDeclaration#[{
Type getDestType() {
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
index fb560acc3dc48..e49b95cee05fa 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
@@ -84,11 +84,6 @@ void populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
// 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 number of contiguous elements grouped per lane by
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index d8eca95cbf23c..150ff0ccb8ade 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -498,7 +498,7 @@ static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr);
+ /*layout=*/nullptr, /*coalesce_hint=*/nullptr);
rewriter.replaceOp(readOp, gatherOp.getResult());
return success();
@@ -533,7 +533,7 @@ static LogicalResult lowerToScatteredStoreOp(vector::TransferWriteOp writeOp,
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr);
+ /*layout=*/nullptr, /*coalesce_hint=*/nullptr);
rewriter.eraseOp(writeOp);
return success();
}
@@ -789,7 +789,7 @@ struct GatherLowering : public OpRewritePattern<vector::GatherOp> {
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr);
+ /*layout=*/nullptr, /*coalesce_hint=*/nullptr);
auto selectOp =
arith::SelectOp::create(rewriter, loc, gatherOp.getMask(),
@@ -824,7 +824,7 @@ struct ScatterLowering : public OpRewritePattern<vector::ScatterOp> {
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr);
+ /*layout=*/nullptr, /*coalesce_hint=*/nullptr);
rewriter.eraseOp(scatterOp);
return success();
}
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index 64be7fca5d40f..ef98ae84523a0 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -113,6 +113,29 @@ isValidGatherScatterBufferParams(Type offsetsTy, Type maskTy,
return success();
}
+// Validates a user-provided (or analysis-stamped) `coalesce_hint` against the
+// op's offsets type. The hint requests grouping `factor` contiguous elements
+// per lane along the innermost (fastest-changing) dimension, so that dimension
+// must be a multiple of `factor`. The further `lane_layout * factor` split is
+// chip-dependent (subgroup size) and is checked when the hint is lowered to a
+// layout, not here.
+static LogicalResult
+isValidCoalesceHint(xegpu::CoalesceHintAttr hint, Type offsetsTy,
+ function_ref<InFlightDiagnostic()> emitError) {
+ if (!hint)
+ return success();
+ auto offsetsVecTy = dyn_cast<VectorType>(offsetsTy);
+ if (!offsetsVecTy)
+ return emitError()
+ << "coalesce_hint requires vector offsets (one per lane).";
+ int64_t factor = hint.getFactor().getInt();
+ int64_t inner = offsetsVecTy.getShape().back();
+ if (inner % factor != 0)
+ return emitError() << "coalesce_hint factor " << factor
+ << " must divide the innermost offsets dim " << inner;
+ return success();
+}
+
LogicalResult
IsValidMatrixOpParams(VectorType dataTy, MemDescType mdescTy,
UnitAttr subgroup_block_io, DistributeLayoutAttr layout,
@@ -582,6 +605,9 @@ LogicalResult LoadGatherOp::verify() {
}
auto offsetsTy = getOffsets().getType();
+ if (failed(isValidCoalesceHint(getCoalesceHintAttr(), offsetsTy,
+ [&]() { return emitOpError(); })))
+ return failure();
return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy, chunkSize,
[&]() { return emitOpError(); });
}
@@ -599,7 +625,7 @@ void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
auto offset = vector::FromElementsOp::create(builder, loc, type, values);
build(builder, state, valueType, source, offset, mask, chunk_size, l1_hint,
- l2_hint, l3_hint, /*anchor_layout=*/nullptr);
+ l2_hint, l3_hint, /*anchor_layout=*/nullptr, /*coalesce_hint=*/nullptr);
}
void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
@@ -616,7 +642,7 @@ void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
auto offset = vector::FromElementsOp::create(builder, loc, type, values);
build(builder, state, valueType, source, offset, mask, chunk_size, l1_hint,
- l2_hint, l3_hint, layout);
+ l2_hint, l3_hint, layout, /*coalesce_hint=*/nullptr);
}
//===----------------------------------------------------------------------===//
@@ -648,6 +674,9 @@ LogicalResult StoreScatterOp::verify() {
}
auto offsetsTy = getOffsets().getType();
+ if (failed(isValidCoalesceHint(getCoalesceHintAttr(), offsetsTy,
+ [&]() { return emitOpError(); })))
+ return failure();
return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy, chunkSize,
[&]() { return emitOpError(); });
}
@@ -667,7 +696,7 @@ void StoreScatterOp::build(OpBuilder &builder, OperationState &state,
// Call the correct builder overload that does not expect result types.
build(builder, state, value, dest, offset, mask, chunk_size, l1_hint, l2_hint,
- l3_hint, /*anchor_layout=*/nullptr);
+ l3_hint, /*anchor_layout=*/nullptr, /*coalesce_hint=*/nullptr);
}
void StoreScatterOp::build(
@@ -683,7 +712,7 @@ void StoreScatterOp::build(
// Call the correct builder overload that does not expect result types.
build(builder, state, value, dest, offset, mask, chunk_size, l1_hint, l2_hint,
- l3_hint, layout);
+ l3_hint, layout, /*coalesce_hint=*/nullptr);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
index 788a5774afc91..89b6c5ac9c7da 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
@@ -1068,6 +1068,13 @@ static void analyzeAndStampHint(OpTy op, DataFlowSolver &solver,
unsigned maxChunkSize) {
if (!isCandidateForCoalesce(op))
return;
+ // A user-provided `coalesce_hint` takes precedence: if the op already
+ // carries one (authored by the user, or stamped by an earlier propagation
+ // level), leave it untouched. The analysis is only a fallback that fills in
+ // a hint where the user did not request one. This also makes the analysis
+ // idempotent across the lane and inst-data propagation runs.
+ if (op.getCoalesceHintAttr())
+ return;
// Do not coalesce accesses tied to a reduction (see isReductionTied);
// reduction coalescing is added in follow-up PRs.
if (isReductionTied(op))
@@ -1081,8 +1088,8 @@ static void analyzeAndStampHint(OpTy op, DataFlowSolver &solver,
decide(lat->getValue(), offsetsTy.getShape(), maxChunkSize, subgroupSize);
if (d.kind != CoalesceDecision::Kind::Chunked)
return;
- auto hint = xegpu::CoalesceHintAttr::get(op.getContext(), d.factor);
- op->setAttr(xegpu::getCoalesceHintAttrName(), hint);
+ op.setCoalesceHintAttr(
+ xegpu::CoalesceHintAttr::get(op.getContext(), d.factor));
}
/// Apply a stamped hint on `op`: build a lane_layout/lane_data/inst_data
@@ -1092,8 +1099,7 @@ static void analyzeAndStampHint(OpTy op, DataFlowSolver &solver,
/// malformed.
template <typename OpTy>
static LogicalResult applyHintOnOp(OpTy op) {
- auto hint = op->template getAttrOfType<xegpu::CoalesceHintAttr>(
- xegpu::getCoalesceHintAttrName());
+ auto hint = op.getCoalesceHintAttr();
if (!hint)
return success(); // no hint: idempotent no-op.
@@ -1113,7 +1119,7 @@ static LogicalResult applyHintOnOp(OpTy op) {
auto layout = buildLaneDataLayout(op.getContext(), valueTy.getRank(),
laneLayout, factor);
op.setLayoutAttr(layout);
- op->removeAttr(xegpu::getCoalesceHintAttrName());
+ op.removeCoalesceHintAttr();
return success();
}
@@ -1126,9 +1132,8 @@ static LogicalResult applyHintOnOp(Operation *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(xegpu::getCoalesceHintAttrName()))
- op->removeAttr(xegpu::getCoalesceHintAttrName());
+ // `coalesce_hint` is an inherent attribute of LoadGatherOp / StoreScatterOp
+ // only, so no other op can carry one.
return success();
}
@@ -1159,16 +1164,14 @@ void mlir::xegpu::runCoalesceGatherScatterAnalysis(
}
void mlir::xegpu::coalesceGatherScatter(Operation *root) {
- root->walk([&](Operation *op) {
- if (op->hasAttr(xegpu::getCoalesceHintAttrName()))
- (void)applyHintOnOp(op);
- });
+ root->walk([&](Operation *op) { (void)applyHintOnOp(op); });
}
void mlir::xegpu::clearCoalesceGatherScatterHints(Operation *root) {
- StringRef name = xegpu::getCoalesceHintAttrName();
root->walk([&](Operation *op) {
- if (op->hasAttr(name))
- op->removeAttr(name);
+ if (auto load = dyn_cast<xegpu::LoadGatherOp>(op))
+ load.removeCoalesceHintAttr();
+ else if (auto store = dyn_cast<xegpu::StoreScatterOp>(op))
+ store.removeCoalesceHintAttr();
});
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 8a926754e7cfb..27ff32666a59d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -503,7 +503,7 @@ struct SgToLaneLoadGather : public OpConversionPattern<xegpu::LoadGatherOp> {
auto newOp = xegpu::LoadGatherOp::create(
rewriter, op.getLoc(), distResultTy1D, distSource, distOffsets,
distMask, op.getChunkSizeAttr(), op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), /*layout=*/nullptr);
+ op.getL3HintAttr(), /*layout=*/nullptr, /*coalesce_hint=*/nullptr);
Value result = newOp->getResult(0);
if (distResultTy1D != distResultTy)
@@ -1034,7 +1034,8 @@ struct SgToLaneStoreScatter
xegpu::StoreScatterOp::create(rewriter, op.getLoc(), distValue, distDest,
distOffsets, distMask, op.getChunkSizeAttr(),
op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), /*layout=*/nullptr);
+ op.getL3HintAttr(), /*layout=*/nullptr,
+ /*coalesce_hint=*/nullptr);
rewriter.eraseOp(op);
return success();
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index aab36b79845e4..292075b0a8394 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -566,7 +566,8 @@ struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
auto newOp = xegpu::LoadGatherOp::create(
rewriter, loc, newValueTy, op.getSource(), o, m,
rewriter.getI64IntegerAttr(chunkSize), op.getL1HintAttr(),
- op.getL2HintAttr(), op.getL3HintAttr(), layout);
+ op.getL2HintAttr(), op.getL3HintAttr(), layout,
+ /*coalesce_hint=*/nullptr);
newOps.push_back(newOp);
}
@@ -661,7 +662,8 @@ struct UnrollStoreScatterOp : public UnrollPattern<xegpu::StoreScatterOp> {
xegpu::StoreScatterOp::create(rewriter, loc, v, op.getDest(), o, m,
rewriter.getI64IntegerAttr(chunkSize),
op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), layout);
+ op.getL3HintAttr(), layout,
+ /*coalesce_hint=*/nullptr);
}
rewriter.eraseOp(op);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 1aa03ebc0f376..4f3a1c2a8c947 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -901,7 +901,7 @@ struct WgToSgLoadGatherOp : public OpConversionPattern<xegpu::LoadGatherOp> {
auto newLoadOp = xegpu::LoadGatherOp::create(
rewriter, loc, newTy, op.getSource(), offsets, mask, chunkSizeAttr,
op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(),
- newLayout);
+ newLayout, /*coalesce_hint=*/nullptr);
newLoadOps.push_back(newLoadOp);
}
rewriter.replaceOpWithMultiple(op, {newLoadOps});
@@ -947,7 +947,8 @@ struct WgToSgStoreScatterOp
xegpu::StoreScatterOp::create(rewriter, loc, val, op.getDest(), offs,
mask, chunkSizeAttr, op.getL1HintAttr(),
op.getL2HintAttr(), op.getL3HintAttr(),
- layout.dropSgLayoutAndData());
+ layout.dropSgLayoutAndData(),
+ /*coalesce_hint=*/nullptr);
}
rewriter.eraseOp(op);
return success();
diff --git a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
index 4b528704607a9..ac1d644dfcd47 100644
--- a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
+++ b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
@@ -1,7 +1,7 @@
// 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
+// Analyze-only mode: stamps the `coalesce_hint` attribute on coalescible ops and
// leaves the layout unchanged. This test pins the hint attribute contract
// that the apply API (and downstream propagator integrations) consume.
@@ -9,7 +9,7 @@
// 1-D vector.step, fully coalescible -> hint with factor = 2 stamped.
// CHECK-LABEL: func.func @load_step_offsets(
// CHECK: xegpu.load
-// CHECK-SAME: {xegpu.coalesce_hint = #xegpu.coalesce_hint<factor = 2 : i64>}
+// CHECK-SAME: <{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> {
@@ -24,7 +24,7 @@ func.func @load_step_offsets(%ptr: i64) -> 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: <{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> {
@@ -74,7 +74,7 @@ func.func @load_broadcast_offsets_no_hint(%ptr: i64) -> vector<32xf32> {
// op.
// CHECK-LABEL: func.func @store_step_hint(
// CHECK: xegpu.store
-// CHECK-SAME: {xegpu.coalesce_hint = #xegpu.coalesce_hint<factor = 2 : i64>}
+// CHECK-SAME: <{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>) {
@@ -84,3 +84,17 @@ func.func @store_step_hint(%ptr: i64, %v: vector<32xf32>) {
: vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
return
}
+
+// -----
+// A user-provided coalesce_hint takes precedence: the analysis must not
+// overwrite it, even if its own heuristic would pick a different factor.
+// CHECK-LABEL: func.func @user_hint_preserved(
+// CHECK: xegpu.load
+// CHECK-SAME: <{coalesce_hint = #xegpu.coalesce_hint<factor = 2 : i64>}>
+func.func @user_hint_preserved(%ptr: i64) -> vector<32xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 2>}>
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index 4f45ce574232a..6e7908836043f 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -772,3 +772,27 @@ func.func @dpas_mx_scale_b_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b
%1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a_val scale_b = %scale_b_val {layout_b_scale = #layout_b_scale_invalid} : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
return
}
+
+// -----
+func.func @coalesce_hint_not_power_of_two(%src: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
+ // expected-error at +1 {{'factor' : 6 must be a power of two}}
+ %val = xegpu.load %src[%offset], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 6>}>
+ : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ return
+}
+
+// -----
+func.func @coalesce_hint_factor_too_small(%src: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
+ // expected-error at +1 {{'factor' : 1 must be >= 2}}
+ %val = xegpu.load %src[%offset], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 1>}>
+ : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ return
+}
+
+// -----
+func.func @coalesce_hint_factor_does_not_divide(%src: i64, %offset: vector<6xindex>, %mask: vector<6xi1>) {
+ // expected-error at +1 {{coalesce_hint factor 4 must divide the innermost offsets dim 6}}
+ %val = xegpu.load %src[%offset], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 4>}>
+ : i64, vector<6xindex>, vector<6xi1> -> vector<6xf32>
+ return
+}
diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index 0b1ac71cdbd32..d0ba3c2119072 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -452,6 +452,23 @@ gpu.func @subgroup_store_offset_1(%dest: memref<?xf16>) {
gpu.return
}
+// CHECK: gpu.func @load_coalesce_hint(%[[arg0:.*]]: i64, %[[arg1:.*]]: vector<16xindex>, %[[arg2:.*]]: vector<16xi1>) {
+gpu.func @load_coalesce_hint(%src: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
+ // A user-provided `coalesce_hint` round-trips through the optional op attribute.
+ // CHECK: xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] <{coalesce_hint = #xegpu.coalesce_hint<factor = 4 : i64>}> : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ %val = xegpu.load %src[%offset], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 4>}>
+ : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ gpu.return
+}
+
+// CHECK: gpu.func @store_coalesce_hint(%[[arg0:.*]]: vector<16xf32>, %[[arg1:.*]]: i64, %[[arg2:.*]]: vector<16xindex>, %[[arg3:.*]]: vector<16xi1>) {
+gpu.func @store_coalesce_hint(%val: vector<16xf32>, %dest: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
+ // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] <{coalesce_hint = #xegpu.coalesce_hint<factor = 4 : i64>}> : vector<16xf32>, i64, vector<16xindex>, vector<16xi1>
+ xegpu.store %val, %dest[%offset], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 4>}>
+ : vector<16xf32>, i64, vector<16xindex>, vector<16xi1>
+ gpu.return
+}
+
// CHECK: gpu.func @prefetch_offset(%[[arg0:.*]]: ui64) {
gpu.func @prefetch_offset(%src: ui64) {
//CHECK: %[[cst:.*]] = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
>From f4e39af355551d7229160a843ba6225cbf676d70 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Thu, 18 Jun 2026 16:04:25 +0000
Subject: [PATCH 5/7] Fix clang-format.
---
mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp | 3 ++-
mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp | 4 ++--
2 files changed, 4 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 150ff0ccb8ade..c9cc06614629b 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -824,7 +824,8 @@ struct ScatterLowering : public OpRewritePattern<vector::ScatterOp> {
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr, /*coalesce_hint=*/nullptr);
+ /*layout=*/nullptr,
+ /*coalesce_hint=*/nullptr);
rewriter.eraseOp(scatterOp);
return success();
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 4f3a1c2a8c947..fa1b598767de4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -900,8 +900,8 @@ struct WgToSgLoadGatherOp : public OpConversionPattern<xegpu::LoadGatherOp> {
auto newLayout = layout.dropSgLayoutAndData();
auto newLoadOp = xegpu::LoadGatherOp::create(
rewriter, loc, newTy, op.getSource(), offsets, mask, chunkSizeAttr,
- op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(),
- newLayout, /*coalesce_hint=*/nullptr);
+ op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(), newLayout,
+ /*coalesce_hint=*/nullptr);
newLoadOps.push_back(newLoadOp);
}
rewriter.replaceOpWithMultiple(op, {newLoadOps});
>From 14bf5ffedd3e97b9548de45068c95bda17184cef Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Wed, 24 Jun 2026 15:14:39 +0000
Subject: [PATCH 6/7] Address review comments.
- rename the analysis file to `XeGPUContiguityAnalysis.cpp`
- keep the file analysis only, and move all mechanism of apply
to test pass.
- reduce comment clutters.
- replace the coalesce_hint attribute by I64Attr `contiguous_chunk`
---
.../mlir/Dialect/XeGPU/IR/XeGPUAttrs.td | 43 --
.../include/mlir/Dialect/XeGPU/IR/XeGPUOps.td | 4 +-
.../Dialect/XeGPU/Transforms/Transforms.h | 34 +-
.../VectorToXeGPU/VectorToXeGPU.cpp | 8 +-
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 15 -
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp | 45 +-
...catter.cpp => XeGPUContiguityAnalysis.cpp} | 404 +++---------------
.../Transforms/XeGPUSgToLaneDistribute.cpp | 4 +-
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 4 +-
.../Transforms/XeGPUWgToSgDistribute.cpp | 4 +-
.../coalesce-gather-scatter-analyze.mlir | 54 ++-
mlir/test/Dialect/XeGPU/invalid.mlir | 20 +-
mlir/test/Dialect/XeGPU/ops.mlir | 18 +-
.../lib/Dialect/XeGPU/TestXeGPUTransforms.cpp | 165 ++++++-
14 files changed, 300 insertions(+), 522 deletions(-)
rename mlir/lib/Dialect/XeGPU/Transforms/{XeGPUCoalesceGatherScatter.cpp => XeGPUContiguityAnalysis.cpp} (68%)
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
index e16ee5031785c..40edce8a60429 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
@@ -927,49 +927,6 @@ def XeGPU_MemLayoutAttr : XeGPUAttr<"MemLayout", "mem_layout"> {
}
-def XeGPU_CoalesceHintAttr : XeGPUAttr<"CoalesceHint", "coalesce_hint"> {
- let summary = [{Optional request to coalesce a gather/scatter access.}];
-
- let description = [{
- `CoalesceHintAttr` is the optional `coalesce_hint` attribute of
- `xegpu.load` / `xegpu.store`. It requests that the access group `factor`
- contiguous elements per lane along the innermost (fastest-changing)
- dimension of the offsets. `factor` must be a power of two `>= 2`, and the
- innermost offsets dimension must be a multiple of it.
-
- The hint has two sources:
- - it can be set directly by the user (or a higher-level frontend that
- already knows the access is contiguous-per-lane), in which case the
- coalesce-gather-scatter analysis is not consulted for that op; or
- - it is stamped by the analysis when it proves the access coalescible.
-
- The hint is a *request*, not a commitment: a consumer turns `factor` into
- a concrete `xegpu.layout` (deriving the `lane_layout` / `lane_data` split
- from the op's offsets inner extent and the chip's subgroup size) and
- clears the hint. A consumer that declines (for example, layout propagation
- detecting a conflict with an anchor-driven layout) likewise clears it
- rather than leaving it dangling.
-
- Example:
- ```mlir
- %v = xegpu.load %ptr[%offsets], %mask <{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/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index a58ee7430a56d..5255448dae33d 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -774,7 +774,7 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l3_hint,
OptionalAttr<DistributeLayoutAttr>:$layout,
- OptionalAttr<XeGPU_CoalesceHintAttr>:$coalesce_hint);
+ OptionalAttr<I64Attr>:$contiguous_chunk);
let results = (outs XeGPU_ValueOrScalarType:$value);
let extraClassDeclaration = extraBaseClassDeclaration # [{
@@ -905,7 +905,7 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l3_hint,
OptionalAttr<DistributeLayoutAttr>:$layout,
- OptionalAttr<XeGPU_CoalesceHintAttr>:$coalesce_hint);
+ OptionalAttr<I64Attr>:$contiguous_chunk);
let extraClassDeclaration = extraBaseClassDeclaration#[{
Type getDestType() {
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
index e49b95cee05fa..84b4dcfb0a736 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
@@ -84,33 +84,13 @@ void populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
// Coalesce gather/scatter analysis + apply.
//===----------------------------------------------------------------------===//
-/// Options controlling `runCoalesceGatherScatterAnalysis`.
-struct CoalesceGatherScatterAnalysisOptions {
- /// Upper bound on the number of contiguous elements grouped per lane 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` or a non-uniform mask are skipped (no hint stamped). The hint
-/// is consumed downstream — either by `coalesceGatherScatter` or by a
-/// layout-propagation pass that reads the hint directly.
-void runCoalesceGatherScatterAnalysis(
- Operation *root, const CoalesceGatherScatterAnalysisOptions &options = {});
-
-/// Walk `root` and turn every stamped `xegpu.coalesce_hint` into an
-/// equivalent `lane_layout` / `lane_data` / `inst_data` layout on its op,
-/// then remove the hint. A hint on an op that cannot be coalesced (e.g. one
-/// that isn't a gather/scatter) is simply dropped. Pairs with
-/// `runCoalesceGatherScatterAnalysis`.
-void coalesceGatherScatter(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);
+/// Run the AxisInfo-based contiguity analysis over `root` and stamp a
+/// `contiguous_chunk` attribute on every `xegpu.load` / `xegpu.store` whose
+/// offsets are contiguous (in runs of >= 2) along the innermost dimension.
+/// The stamped value is the inner-dim contiguity; it is a target-independent
+/// property consumed downstream (e.g. to derive a `lane_data` split). Ops that
+/// already carry a `contiguous_chunk` attribute are left untouched.
+void runCoalesceGatherScatterAnalysis(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
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index c9cc06614629b..a428a9d82e618 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -498,7 +498,7 @@ static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr, /*coalesce_hint=*/nullptr);
+ /*layout=*/nullptr, /*contiguous_chunk=*/nullptr);
rewriter.replaceOp(readOp, gatherOp.getResult());
return success();
@@ -533,7 +533,7 @@ static LogicalResult lowerToScatteredStoreOp(vector::TransferWriteOp writeOp,
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr, /*coalesce_hint=*/nullptr);
+ /*layout=*/nullptr, /*contiguous_chunk=*/nullptr);
rewriter.eraseOp(writeOp);
return success();
}
@@ -789,7 +789,7 @@ struct GatherLowering : public OpRewritePattern<vector::GatherOp> {
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr, /*coalesce_hint=*/nullptr);
+ /*layout=*/nullptr, /*contiguous_chunk=*/nullptr);
auto selectOp =
arith::SelectOp::create(rewriter, loc, gatherOp.getMask(),
@@ -825,7 +825,7 @@ struct ScatterLowering : public OpRewritePattern<vector::ScatterOp> {
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
/*layout=*/nullptr,
- /*coalesce_hint=*/nullptr);
+ /*contiguous_chunk=*/nullptr);
rewriter.eraseOp(scatterOp);
return success();
}
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 8cdef2dd994d8..e92b109c2223e 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -1213,21 +1213,6 @@ 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/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index ef98ae84523a0..1ef90f5915934 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -113,26 +113,26 @@ isValidGatherScatterBufferParams(Type offsetsTy, Type maskTy,
return success();
}
-// Validates a user-provided (or analysis-stamped) `coalesce_hint` against the
-// op's offsets type. The hint requests grouping `factor` contiguous elements
-// per lane along the innermost (fastest-changing) dimension, so that dimension
-// must be a multiple of `factor`. The further `lane_layout * factor` split is
-// chip-dependent (subgroup size) and is checked when the hint is lowered to a
-// layout, not here.
+// Validates the `contiguous_chunk` attribute against the op's offsets type:
+// the inner offsets dimension is contiguous in runs of `size`, so `size` must
+// be in [2, innermost offsets dim].
static LogicalResult
-isValidCoalesceHint(xegpu::CoalesceHintAttr hint, Type offsetsTy,
- function_ref<InFlightDiagnostic()> emitError) {
- if (!hint)
+isValidContiguousChunk(std::optional<uint64_t> chunk, Type offsetsTy,
+ function_ref<InFlightDiagnostic()> emitError) {
+ if (!chunk)
return success();
auto offsetsVecTy = dyn_cast<VectorType>(offsetsTy);
if (!offsetsVecTy)
return emitError()
- << "coalesce_hint requires vector offsets (one per lane).";
- int64_t factor = hint.getFactor().getInt();
+ << "contiguous_chunk requires vector offsets (one per lane).";
+ int64_t size = static_cast<int64_t>(*chunk);
int64_t inner = offsetsVecTy.getShape().back();
- if (inner % factor != 0)
- return emitError() << "coalesce_hint factor " << factor
- << " must divide the innermost offsets dim " << inner;
+ if (size < 2)
+ return emitError() << "contiguous_chunk " << size << " must be >= 2";
+ if (size > inner)
+ return emitError() << "contiguous_chunk " << size
+ << " must not exceed the innermost offsets dim "
+ << inner;
return success();
}
@@ -605,8 +605,8 @@ LogicalResult LoadGatherOp::verify() {
}
auto offsetsTy = getOffsets().getType();
- if (failed(isValidCoalesceHint(getCoalesceHintAttr(), offsetsTy,
- [&]() { return emitOpError(); })))
+ if (failed(isValidContiguousChunk(getContiguousChunk(), offsetsTy,
+ [&]() { return emitOpError(); })))
return failure();
return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy, chunkSize,
[&]() { return emitOpError(); });
@@ -625,7 +625,8 @@ void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
auto offset = vector::FromElementsOp::create(builder, loc, type, values);
build(builder, state, valueType, source, offset, mask, chunk_size, l1_hint,
- l2_hint, l3_hint, /*anchor_layout=*/nullptr, /*coalesce_hint=*/nullptr);
+ l2_hint, l3_hint, /*anchor_layout=*/nullptr,
+ /*contiguous_chunk=*/nullptr);
}
void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
@@ -642,7 +643,7 @@ void LoadGatherOp::build(OpBuilder &builder, OperationState &state,
auto offset = vector::FromElementsOp::create(builder, loc, type, values);
build(builder, state, valueType, source, offset, mask, chunk_size, l1_hint,
- l2_hint, l3_hint, layout, /*coalesce_hint=*/nullptr);
+ l2_hint, l3_hint, layout, /*contiguous_chunk=*/nullptr);
}
//===----------------------------------------------------------------------===//
@@ -674,8 +675,8 @@ LogicalResult StoreScatterOp::verify() {
}
auto offsetsTy = getOffsets().getType();
- if (failed(isValidCoalesceHint(getCoalesceHintAttr(), offsetsTy,
- [&]() { return emitOpError(); })))
+ if (failed(isValidContiguousChunk(getContiguousChunk(), offsetsTy,
+ [&]() { return emitOpError(); })))
return failure();
return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy, chunkSize,
[&]() { return emitOpError(); });
@@ -696,7 +697,7 @@ void StoreScatterOp::build(OpBuilder &builder, OperationState &state,
// Call the correct builder overload that does not expect result types.
build(builder, state, value, dest, offset, mask, chunk_size, l1_hint, l2_hint,
- l3_hint, /*anchor_layout=*/nullptr, /*coalesce_hint=*/nullptr);
+ l3_hint, /*anchor_layout=*/nullptr, /*contiguous_chunk=*/nullptr);
}
void StoreScatterOp::build(
@@ -712,7 +713,7 @@ void StoreScatterOp::build(
// Call the correct builder overload that does not expect result types.
build(builder, state, value, dest, offset, mask, chunk_size, l1_hint, l2_hint,
- l3_hint, layout, /*coalesce_hint=*/nullptr);
+ l3_hint, layout, /*contiguous_chunk=*/nullptr);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp
similarity index 68%
rename from mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
rename to mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp
index 89b6c5ac9c7da..1c67503f7445a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCoalesceGatherScatter.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp
@@ -6,31 +6,19 @@
//
//===----------------------------------------------------------------------===//
//
-// This file implements the gather/scatter coalescing analysis. It decides
-// whether an `xegpu.load` / `xegpu.store` gather/scatter accesses `N`
-// contiguous elements per lane along the innermost dimension, and, if so,
-// stamps a `xegpu.coalesce_hint<factor = N>` attribute recording the chosen
-// factor. The decision is driven by a small XeGPU-local axis-info dataflow
-// analysis tracking per-axis `contiguity`, `constancy`, and `divisibility`.
+// This file implements the contiguity analysis. It computes, for a memory
+// operation (i.e., `xegpu.load` / `xegpu.store`), how many elements are
+// contiguous along the innermost offsets dimension, and stamps that count as an
+// attribute on the op. The analysis is a small XeGPU-local
+// axis-info dataflow tracking per-axis `contiguity`, `constancy`, and
+// `divisibility`; the stamped value is the inner-dim `contiguity`.
//
-// The analysis performs no rewrite. The hint it stamps is turned into a
-// `lane_data` layout by `coalesceGatherScatter` (used by the test pass) or
-// read directly by layout propagation; the actual memory-message rewrite is
-// left to the downstream WG-to-SG / SG-to-Lane distribution passes, which
-// interpret `lane_data`.
+// Contiguity is a target-independent property of the offsets.
//
// 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.
-//
-// 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.
+// type at any rank, against the innermost dimension.
//
+// The analysis gets it's inspiration from the Triton Axis info analysis.
//===----------------------------------------------------------------------===//
#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
@@ -43,14 +31,9 @@
#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"
#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/ADT/bit.h"
#include "llvm/Support/MathExtras.h"
#include <numeric>
#include <optional>
@@ -74,25 +57,16 @@ namespace mlir::xegpu::detail::axis_dataflow {
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.
-/// - `innerStride`: when set, every "row" along the innermost dimension
-/// is an arithmetic progression with this (constant) step. This is the
-/// general form of which `contiguity`/`constancy` are the two special
-/// cases: `innerStride = 1` is the stride-1 contiguous case (and implies
-/// `contiguity[innerDim] > 1`), `innerStride = 0` is the all-equal case
-/// (implies `constancy[innerDim] > 1`), and any other value (e.g. 4) is a
-/// strided progression that neither `contiguity` nor `constancy` can
-/// represent — both read 1 there. Per-row base may differ across outer
-/// indices; per-row alignment is captured by `divisibility[innerDim]`.
-/// The coalescing decision only acts on the `innerStride = 1` case.
+/// All fields describe runs of consecutive elements along a dimension `d`:
+/// - `contiguity[d]`: longest run that increases by exactly 1.
+/// - `constancy[d]`: longest run of equal values.
+/// - `divisibility[d]`: a power-of-two divisor of every element.
+/// - `knownConstant`: the value, if the whole vector is one constant.
+/// - `innerStride`: if set, each row along the innermost dim is an
+/// arithmetic progression with this step. `1` is the contiguous case,
+/// `0` the all-equal case; any other value is a strided progression that
+/// `contiguity`/`constancy` can't represent (both read 1). Per-row base
+/// may vary; per-row alignment lives in `divisibility[innerDim]`.
///
/// Pessimistic / entry value: contiguity=1, constancy=1, divisibility=1,
/// innerStride absent.
@@ -586,9 +560,9 @@ class AxisInfoAnalysis
unsigned r = vt.getRank();
auto shape = vt.getShape();
AxisInfo v = AxisInfo::getPessimistic(r);
- auto unitConstant = [](const AxisInfo &a, unsigned d, int64_t lanes) {
+ auto unitConstant = [](const AxisInfo &a, unsigned d, int64_t extent) {
return a.knownConstant && *a.knownConstant == 1 &&
- a.constancy[d] >= lanes;
+ a.constancy[d] >= extent;
};
for (unsigned d = 0; d < r; ++d) {
v.constancy[d] = std::min({shape[d], lhs.constancy[d], rhs.constancy[d]});
@@ -617,22 +591,17 @@ class AxisInfoAnalysis
}
// 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.
+ // positive constant `c`. The lhs must be an arithmetic progression (AP)
+ // along the inner dim, i.e. its values step by a constant stride `s`, and
+ // `c` must divide `s`.
//
- // 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).
+ // Take the inner row `[0, 2, 4, 6, 8, 10, 12, 14]` (stride s = 2) and c = 2:
+ // - Division `/ 2` gives `[0, 1, 2, 3, 4, 5, 6, 7]`: a new AP with stride
+ // `s / c = 1`. A resulting stride of 1 is contiguous, 0 is constant.
+ // - Remainder `% 2` gives `[0, 0, 0, 0, 0, 0, 0, 0]`: every element folds
+ // to the same residue, so the row is constant (stride 0).
//
- // Signed vs unsigned only differs in the constant interpretation; we
- // require positive constants so the signed/unsigned distinction is moot
- // here.
+ // We require positive `c`, so signed and unsigned behave the same.
template <bool IsSigned, bool IsRem, typename OpTy>
LogicalResult visitDivRem(OpTy op, ArrayRef<const AxisInfoLattice *> operands,
ArrayRef<AxisInfoLattice *> results) {
@@ -694,12 +663,16 @@ class AxisInfoAnalysis
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.
+ // arith.andi: `x & m` with a uniform positive constant mask `m`. The
+ // interesting case is `m = P - 1` for a power of 2 `P`, which is the same
+ // as `x % P` (see visitDivRem): masking an inner row whose stride is a
+ // multiple of `P` folds it to a constant.
+ //
+ // Take the row `[0, 2, 4, 6, 8, 10, 12, 14]` (stride 2) and m = 1 (P = 2):
+ // `x & 1` gives `[0, 0, 0, 0, 0, 0, 0, 0]`: constant along the inner dim.
+ //
+ // Also handles the trivial masks `m == 0` (always zero) and all-ones
+ // (identity).
LogicalResult visitAndI(arith::AndIOp op,
ArrayRef<const AxisInfoLattice *> operands,
ArrayRef<AxisInfoLattice *> results) {
@@ -865,286 +838,40 @@ using ::mlir::xegpu::detail::axis_dataflow::AxisInfo;
using ::mlir::xegpu::detail::axis_dataflow::AxisInfoLattice;
//===----------------------------------------------------------------------===//
-// Coalescing decision.
+// Analysis driver.
//===----------------------------------------------------------------------===//
-struct CoalesceDecision {
- 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
-};
-
-/// 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.
-///
-/// 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
-/// `maxChunkSize` and the offsets contiguity.
-static CoalesceDecision decide(const AxisInfo &info,
- ArrayRef<int64_t> offsetsShape,
- unsigned maxChunkSize, unsigned subgroupSize) {
- CoalesceDecision d;
- if (!info.isInitialized() || offsetsShape.empty())
- return d;
- unsigned innerDim = offsetsShape.size() - 1;
- int64_t inner = offsetsShape[innerDim];
- if (inner < 2)
- return d;
-
- // 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.
- 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.
-
- int64_t budget = static_cast<int64_t>(maxChunkSize);
- 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(perLane, bound);
- if (factor < 2)
- return d;
- d.kind = CoalesceDecision::Kind::Chunked;
- d.laneLayout = laneLayout;
- d.factor = factor;
- return d;
-}
-
-/// 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();
-}
-
-/// 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 innerLaneLayout,
- int64_t innerLaneData) {
- SmallVector<int32_t> laneLayout(rank, 1);
- SmallVector<int32_t> laneData(rank, 1);
- 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);
-}
-
-//===----------------------------------------------------------------------===//
-// Analysis driver + hint apply.
-//===----------------------------------------------------------------------===//
-
-/// 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;
-}
-
-/// Common analysis preconditions: vector offsets/value, all-true mask,
-/// no existing non-trivial lane_data.
+/// Stamp a `contiguous_chunk` attribute on `op` recording the inner-dim
+/// contiguity computed by the analysis. The contiguity is a target-independent
+/// property of the offsets.
template <typename OpTy>
-static bool isCandidateForCoalesce(OpTy op) {
+static void analyzeAndStampContiguity(OpTy op, DataFlowSolver &solver) {
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;
- return true;
-}
-
-/// True when `op` (a gather load or scatter store) is tied to a
-/// `vector.multi_reduction`: a load whose result feeds a reduction, or a
-/// store whose stored value comes from one (through layout-neutral /
-/// elementwise / insert glue).
-///
-/// Coalescing such an access is gated off for now: the analysis only sets
-/// `lane_data[FCD]`, and consuming that on a reduction requires reduction-
-/// aware layout handling that is not part of this change. A coalesced store
-/// fed (transitively) by a reduction would seed `lane_data[FCD] = N` that the
-/// reduction result (kept at `lane_data = 1`) cannot match, producing an
-/// unlowerable `xegpu.convert_layout lane_data=[1]<->[N]`. Follow-up PRs that
-/// add reduction coalescing relax this gate.
-template <typename OpTy>
-static bool isReductionTied(OpTy op) {
- if constexpr (std::is_same_v<OpTy, xegpu::StoreScatterOp>) {
- // Backward walk of the stored value's slice through the reassembly
- // (convert_layout / shape_cast / insert(_strided_slice) / elementwise)
- // that typically sits between a reduction and the store.
- SmallVector<Value, 8> worklist{op.getValue()};
- llvm::SmallPtrSet<Operation *, 16> seen;
- unsigned steps = 0;
- while (!worklist.empty() && steps++ < 64) {
- Value v = worklist.pop_back_val();
- Operation *def = v.getDefiningOp();
- if (!def || !seen.insert(def).second)
- continue;
- if (isa<vector::MultiDimReductionOp>(def))
- return true;
- if (isa<vector::ShapeCastOp, vector::BitCastOp, xegpu::ConvertLayoutOp,
- vector::InsertOp, vector::InsertStridedSliceOp>(def) ||
- OpTrait::hasElementwiseMappableTraits(def))
- for (Value operand : def->getOperands())
- if (isa<VectorType>(operand.getType()))
- worklist.push_back(operand);
- }
- return false;
- } else {
- for (Operation *user : op->getUsers()) {
- Operation *u = user;
- while (
- u &&
- isa<vector::ShapeCastOp, vector::BitCastOp, xegpu::ConvertLayoutOp>(
- u)) {
- if (u->getNumResults() != 1 || u->getResult(0).use_empty())
- break;
- u = *u->getResult(0).getUsers().begin();
- }
- if (u && isa<vector::MultiDimReductionOp>(u))
- return true;
- }
- return false;
- }
-}
-
-/// 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;
- // A user-provided `coalesce_hint` takes precedence: if the op already
- // carries one (authored by the user, or stamped by an earlier propagation
- // level), leave it untouched. The analysis is only a fallback that fills in
- // a hint where the user did not request one. This also makes the analysis
- // idempotent across the lane and inst-data propagation runs.
- if (op.getCoalesceHintAttr())
return;
- // Do not coalesce accesses tied to a reduction (see isReductionTied);
- // reduction coalescing is added in follow-up PRs.
- if (isReductionTied(op))
+ // A pre-existing `contiguous_chunk` (user-authored, or stamped by an earlier
+ // run) takes precedence; leave it untouched so the analysis is idempotent.
+ if (op.getContiguousChunk())
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;
- auto d =
- decide(lat->getValue(), offsetsTy.getShape(), maxChunkSize, subgroupSize);
- if (d.kind != CoalesceDecision::Kind::Chunked)
+ const AxisInfo &info = lat->getValue();
+ unsigned innerDim = offsetsTy.getRank() - 1;
+ int64_t inner = offsetsTy.getShape()[innerDim];
+ int64_t chunk = std::min<int64_t>(info.contiguity[innerDim], inner);
+ if (chunk < 2)
return;
- op.setCoalesceHintAttr(
- xegpu::CoalesceHintAttr::get(op.getContext(), d.factor));
-}
-
-/// 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, 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.getCoalesceHintAttr();
- if (!hint)
- return success(); // no hint: idempotent no-op.
-
- 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);
- op.setLayoutAttr(layout);
- op.removeCoalesceHintAttr();
- return success();
-}
-
-/// Apply a stamped `xegpu.coalesce_hint` on a single op. A hint on an op
-/// that cannot be coalesced (anything other than a gather/scatter) is just
-/// dropped. Returns failure only when the hint is malformed for a real
-/// gather/scatter.
-static LogicalResult applyHintOnOp(Operation *op) {
- if (auto load = dyn_cast<xegpu::LoadGatherOp>(op))
- return applyHintOnOp(load);
- if (auto store = dyn_cast<xegpu::StoreScatterOp>(op))
- return applyHintOnOp(store);
- // `coalesce_hint` is an inherent attribute of LoadGatherOp / StoreScatterOp
- // only, so no other op can carry one.
- return success();
+ op.setContiguousChunk(chunk);
}
} // namespace
//===----------------------------------------------------------------------===//
-// Public APIs.
+// Public API.
//===----------------------------------------------------------------------===//
-void mlir::xegpu::runCoalesceGatherScatterAnalysis(
- Operation *root, const CoalesceGatherScatterAnalysisOptions &options) {
+void mlir::xegpu::runCoalesceGatherScatterAnalysis(Operation *root) {
DataFlowSolver solver;
solver.load<dataflow::DeadCodeAnalysis>();
solver.load<mlir::xegpu::detail::axis_dataflow::AxisInfoAnalysis>();
@@ -1154,24 +881,11 @@ void mlir::xegpu::runCoalesceGatherScatterAnalysis(
// The solver computed AxisInfo for the whole region in the single
// `initializeAndRun` above; offsets shared by several gather/scatter ops are
// analyzed only once. This walk is just per-op point lookups into that
- // result (no re-analysis), turning each cached fact into a hint.
- 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);
- });
-}
-
-void mlir::xegpu::coalesceGatherScatter(Operation *root) {
- root->walk([&](Operation *op) { (void)applyHintOnOp(op); });
-}
-
-void mlir::xegpu::clearCoalesceGatherScatterHints(Operation *root) {
+ // result (no re-analysis), turning each cached fact into an attribute.
root->walk([&](Operation *op) {
if (auto load = dyn_cast<xegpu::LoadGatherOp>(op))
- load.removeCoalesceHintAttr();
+ analyzeAndStampContiguity(load, solver);
else if (auto store = dyn_cast<xegpu::StoreScatterOp>(op))
- store.removeCoalesceHintAttr();
+ analyzeAndStampContiguity(store, solver);
});
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 27ff32666a59d..994a7474ceab2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -503,7 +503,7 @@ struct SgToLaneLoadGather : public OpConversionPattern<xegpu::LoadGatherOp> {
auto newOp = xegpu::LoadGatherOp::create(
rewriter, op.getLoc(), distResultTy1D, distSource, distOffsets,
distMask, op.getChunkSizeAttr(), op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), /*layout=*/nullptr, /*coalesce_hint=*/nullptr);
+ op.getL3HintAttr(), /*layout=*/nullptr, /*contiguous_chunk=*/nullptr);
Value result = newOp->getResult(0);
if (distResultTy1D != distResultTy)
@@ -1035,7 +1035,7 @@ struct SgToLaneStoreScatter
distOffsets, distMask, op.getChunkSizeAttr(),
op.getL1HintAttr(), op.getL2HintAttr(),
op.getL3HintAttr(), /*layout=*/nullptr,
- /*coalesce_hint=*/nullptr);
+ /*contiguous_chunk=*/nullptr);
rewriter.eraseOp(op);
return success();
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index 292075b0a8394..558784befe7cd 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -567,7 +567,7 @@ struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
rewriter, loc, newValueTy, op.getSource(), o, m,
rewriter.getI64IntegerAttr(chunkSize), op.getL1HintAttr(),
op.getL2HintAttr(), op.getL3HintAttr(), layout,
- /*coalesce_hint=*/nullptr);
+ /*contiguous_chunk=*/nullptr);
newOps.push_back(newOp);
}
@@ -663,7 +663,7 @@ struct UnrollStoreScatterOp : public UnrollPattern<xegpu::StoreScatterOp> {
rewriter.getI64IntegerAttr(chunkSize),
op.getL1HintAttr(), op.getL2HintAttr(),
op.getL3HintAttr(), layout,
- /*coalesce_hint=*/nullptr);
+ /*contiguous_chunk=*/nullptr);
}
rewriter.eraseOp(op);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index fa1b598767de4..cd68e7f1808c4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -901,7 +901,7 @@ struct WgToSgLoadGatherOp : public OpConversionPattern<xegpu::LoadGatherOp> {
auto newLoadOp = xegpu::LoadGatherOp::create(
rewriter, loc, newTy, op.getSource(), offsets, mask, chunkSizeAttr,
op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(), newLayout,
- /*coalesce_hint=*/nullptr);
+ /*contiguous_chunk=*/nullptr);
newLoadOps.push_back(newLoadOp);
}
rewriter.replaceOpWithMultiple(op, {newLoadOps});
@@ -948,7 +948,7 @@ struct WgToSgStoreScatterOp
mask, chunkSizeAttr, op.getL1HintAttr(),
op.getL2HintAttr(), op.getL3HintAttr(),
layout.dropSgLayoutAndData(),
- /*coalesce_hint=*/nullptr);
+ /*contiguous_chunk=*/nullptr);
}
rewriter.eraseOp(op);
return success();
diff --git a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
index ac1d644dfcd47..8901231e45ff3 100644
--- a/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
+++ b/mlir/test/Dialect/XeGPU/coalesce-gather-scatter-analyze.mlir
@@ -1,15 +1,15 @@
// RUN: mlir-opt -split-input-file \
// RUN: -test-xegpu-coalesce-gather-scatter="analyze-only=true" %s | FileCheck %s
-// Analyze-only mode: stamps the `coalesce_hint` attribute on coalescible ops and
-// leaves the layout unchanged. This test pins the hint attribute contract
-// that the apply API (and downstream propagator integrations) consume.
+// Analyze-only mode: stamps the `contiguous_chunk` attribute on contiguous ops
+// and leaves the layout unchanged. This test pins the attribute contract that
+// the coalescing consumer (and downstream integrations) read.
// -----
-// 1-D vector.step, fully coalescible -> hint with factor = 2 stamped.
+// 1-D vector.step, fully contiguous -> contiguous_chunk = 32 (the inner extent).
// CHECK-LABEL: func.func @load_step_offsets(
// CHECK: xegpu.load
-// CHECK-SAME: <{coalesce_hint = #xegpu.coalesce_hint<factor = 2 : i64>}>
+// CHECK-SAME: <{contiguous_chunk = 32 : i64}>
// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
// CHECK-NOT: lane_data
func.func @load_step_offsets(%ptr: i64) -> vector<32xf32> {
@@ -21,10 +21,10 @@ func.func @load_step_offsets(%ptr: i64) -> vector<32xf32> {
}
// -----
-// 2-D leading-1 dim: hint stamped on the load with factor = 2.
+// 2-D leading-1 dim: contiguity computed on the inner dim -> contiguous_chunk = 32.
// CHECK-LABEL: func.func @load_2d_leading_unit(
// CHECK: xegpu.load
-// CHECK-SAME: <{coalesce_hint = #xegpu.coalesce_hint<factor = 2 : i64>}>
+// CHECK-SAME: <{contiguous_chunk = 32 : i64}>
// CHECK-SAME: : i64, vector<1x32xindex>, vector<1x32xi1> -> vector<1x32xf32>
// CHECK-NOT: lane_data
func.func @load_2d_leading_unit(%ptr: i64) -> vector<1x32xf32> {
@@ -37,12 +37,12 @@ func.func @load_2d_leading_unit(%ptr: i64) -> vector<1x32xf32> {
}
// -----
-// Stride-4 offsets: not coalescible, no hint stamped.
-// CHECK-LABEL: func.func @load_stride4_no_hint(
+// Stride-4 offsets: not contiguous, no attribute stamped.
+// CHECK-LABEL: func.func @load_stride4_no_chunk(
// CHECK: xegpu.load
-// CHECK-NOT: coalesce_hint
+// CHECK-NOT: contiguous_chunk
// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
-func.func @load_stride4_no_hint(%ptr: i64) -> vector<32xf32> {
+func.func @load_stride4_no_chunk(%ptr: i64) -> vector<32xf32> {
%c4 = arith.constant 4 : index
%step = vector.step : vector<32xindex>
%splat = vector.broadcast %c4 : index to vector<32xindex>
@@ -54,14 +54,13 @@ func.func @load_stride4_no_hint(%ptr: i64) -> 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(
+// All-equal offsets: inner dim is constant (not contiguous), so no attribute
+// is stamped.
+// CHECK-LABEL: func.func @load_broadcast_offsets_no_chunk(
// CHECK: xegpu.load
-// CHECK-NOT: coalesce_hint
+// CHECK-NOT: contiguous_chunk
// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
-func.func @load_broadcast_offsets_no_hint(%ptr: i64) -> vector<32xf32> {
+func.func @load_broadcast_offsets_no_chunk(%ptr: i64) -> vector<32xf32> {
%offsets = arith.constant dense<0> : vector<32xindex>
%mask = arith.constant dense<true> : vector<32xi1>
%v = xegpu.load %ptr[%offsets], %mask
@@ -70,14 +69,13 @@ func.func @load_broadcast_offsets_no_hint(%ptr: i64) -> vector<32xf32> {
}
// -----
-// Store with vector.step offsets: hint stamped with factor = 2 on the store
-// op.
-// CHECK-LABEL: func.func @store_step_hint(
+// Store with vector.step offsets: contiguous_chunk = 32 stamped on the store.
+// CHECK-LABEL: func.func @store_step_chunk(
// CHECK: xegpu.store
-// CHECK-SAME: <{coalesce_hint = #xegpu.coalesce_hint<factor = 2 : i64>}>
+// CHECK-SAME: <{contiguous_chunk = 32 : i64}>
// CHECK-SAME: : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
// CHECK-NOT: lane_data
-func.func @store_step_hint(%ptr: i64, %v: vector<32xf32>) {
+func.func @store_step_chunk(%ptr: i64, %v: vector<32xf32>) {
%offsets = vector.step : vector<32xindex>
%mask = arith.constant dense<true> : vector<32xi1>
xegpu.store %v, %ptr[%offsets], %mask
@@ -86,15 +84,15 @@ func.func @store_step_hint(%ptr: i64, %v: vector<32xf32>) {
}
// -----
-// A user-provided coalesce_hint takes precedence: the analysis must not
-// overwrite it, even if its own heuristic would pick a different factor.
-// CHECK-LABEL: func.func @user_hint_preserved(
+// A pre-existing contiguous_chunk takes precedence: the analysis must not
+// overwrite it.
+// CHECK-LABEL: func.func @user_chunk_preserved(
// CHECK: xegpu.load
-// CHECK-SAME: <{coalesce_hint = #xegpu.coalesce_hint<factor = 2 : i64>}>
-func.func @user_hint_preserved(%ptr: i64) -> vector<32xf32> {
+// CHECK-SAME: <{contiguous_chunk = 2 : i64}>
+func.func @user_chunk_preserved(%ptr: i64) -> vector<32xf32> {
%offsets = vector.step : vector<32xindex>
%mask = arith.constant dense<true> : vector<32xi1>
- %v = xegpu.load %ptr[%offsets], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 2>}>
+ %v = xegpu.load %ptr[%offsets], %mask <{contiguous_chunk = 2 : i64}>
: i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
return %v : vector<32xf32>
}
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index 6e7908836043f..c197788deeb38 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -774,25 +774,17 @@ func.func @dpas_mx_scale_b_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b
}
// -----
-func.func @coalesce_hint_not_power_of_two(%src: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
- // expected-error at +1 {{'factor' : 6 must be a power of two}}
- %val = xegpu.load %src[%offset], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 6>}>
+func.func @contiguous_chunk_too_small(%src: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
+ // expected-error at +1 {{contiguous_chunk 1 must be >= 2}}
+ %val = xegpu.load %src[%offset], %mask <{contiguous_chunk = 1 : i64}>
: i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
return
}
// -----
-func.func @coalesce_hint_factor_too_small(%src: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
- // expected-error at +1 {{'factor' : 1 must be >= 2}}
- %val = xegpu.load %src[%offset], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 1>}>
- : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
- return
-}
-
-// -----
-func.func @coalesce_hint_factor_does_not_divide(%src: i64, %offset: vector<6xindex>, %mask: vector<6xi1>) {
- // expected-error at +1 {{coalesce_hint factor 4 must divide the innermost offsets dim 6}}
- %val = xegpu.load %src[%offset], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 4>}>
+func.func @contiguous_chunk_exceeds_inner(%src: i64, %offset: vector<6xindex>, %mask: vector<6xi1>) {
+ // expected-error at +1 {{contiguous_chunk 8 must not exceed the innermost offsets dim 6}}
+ %val = xegpu.load %src[%offset], %mask <{contiguous_chunk = 8 : i64}>
: i64, vector<6xindex>, vector<6xi1> -> vector<6xf32>
return
}
diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index d0ba3c2119072..7adbbe7dea2ba 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -452,19 +452,19 @@ gpu.func @subgroup_store_offset_1(%dest: memref<?xf16>) {
gpu.return
}
-// CHECK: gpu.func @load_coalesce_hint(%[[arg0:.*]]: i64, %[[arg1:.*]]: vector<16xindex>, %[[arg2:.*]]: vector<16xi1>) {
-gpu.func @load_coalesce_hint(%src: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
- // A user-provided `coalesce_hint` round-trips through the optional op attribute.
- // CHECK: xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] <{coalesce_hint = #xegpu.coalesce_hint<factor = 4 : i64>}> : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
- %val = xegpu.load %src[%offset], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 4>}>
+// CHECK: gpu.func @load_contiguous_chunk(%[[arg0:.*]]: i64, %[[arg1:.*]]: vector<16xindex>, %[[arg2:.*]]: vector<16xi1>) {
+gpu.func @load_contiguous_chunk(%src: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
+ // A user-provided `contiguous_chunk` round-trips through the optional op attribute.
+ // CHECK: xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] <{contiguous_chunk = 4 : i64}> : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ %val = xegpu.load %src[%offset], %mask <{contiguous_chunk = 4 : i64}>
: i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
gpu.return
}
-// CHECK: gpu.func @store_coalesce_hint(%[[arg0:.*]]: vector<16xf32>, %[[arg1:.*]]: i64, %[[arg2:.*]]: vector<16xindex>, %[[arg3:.*]]: vector<16xi1>) {
-gpu.func @store_coalesce_hint(%val: vector<16xf32>, %dest: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
- // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] <{coalesce_hint = #xegpu.coalesce_hint<factor = 4 : i64>}> : vector<16xf32>, i64, vector<16xindex>, vector<16xi1>
- xegpu.store %val, %dest[%offset], %mask <{coalesce_hint = #xegpu.coalesce_hint<factor = 4>}>
+// CHECK: gpu.func @store_contiguous_chunk(%[[arg0:.*]]: vector<16xf32>, %[[arg1:.*]]: i64, %[[arg2:.*]]: vector<16xindex>, %[[arg3:.*]]: vector<16xi1>) {
+gpu.func @store_contiguous_chunk(%val: vector<16xf32>, %dest: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
+ // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] <{contiguous_chunk = 4 : i64}> : vector<16xf32>, i64, vector<16xindex>, vector<16xi1>
+ xegpu.store %val, %dest[%offset], %mask <{contiguous_chunk = 4 : i64}>
: vector<16xf32>, i64, vector<16xindex>, vector<16xi1>
gpu.return
}
diff --git a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
index f24b4dccaab04..f7b768db36460 100644
--- a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
+++ b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
@@ -16,12 +16,16 @@
#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"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Value.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/bit.h"
+#include "llvm/Support/MathExtras.h"
#include "llvm/Support/raw_ostream.h"
#include <optional>
@@ -455,7 +459,7 @@ struct TestXeGPUCoalesceGatherScatter
}
StringRef getDescription() const final {
- return "Test the XeGPU coalesce-gather-scatter analysis + apply APIs.";
+ return "Test the XeGPU contiguity analysis and its coalescing consumer.";
}
void getDependentDialects(::mlir::DialectRegistry ®istry) const override {
@@ -475,16 +479,163 @@ struct TestXeGPUCoalesceGatherScatter
Option<bool> analyzeOnly{
*this, "analyze-only",
- llvm::cl::desc("Only run the analysis (stamp xegpu.coalesce_hint "
+ llvm::cl::desc("Only run the analysis (stamp contiguous_chunk "
"attributes); do not apply."),
llvm::cl::init(false)};
void runOnOperation() override {
- xegpu::CoalesceGatherScatterAnalysisOptions options;
- options.maxChunkSize = maxChunkSize;
- xegpu::runCoalesceGatherScatterAnalysis(getOperation(), options);
- if (!analyzeOnly)
- xegpu::coalesceGatherScatter(getOperation());
+ xegpu::runCoalesceGatherScatterAnalysis(getOperation());
+ if (analyzeOnly)
+ return;
+ getOperation()->walk([&](Operation *op) {
+ if (auto load = dyn_cast<xegpu::LoadGatherOp>(op))
+ applyContiguousChunk(load, maxChunkSize);
+ else if (auto store = dyn_cast<xegpu::StoreScatterOp>(op))
+ applyContiguousChunk(store, maxChunkSize);
+ });
+ }
+
+private:
+ /// 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;
+ }
+
+ /// 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.
+ static unsigned lookupSubgroupSize(Operation *op) {
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
+ return uArch ? static_cast<unsigned>(uArch->getSubgroupSize()) : 16u;
+ }
+
+ /// 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.
+ static xegpu::LayoutAttr buildLaneDataLayout(MLIRContext *ctx, unsigned rank,
+ int64_t innerLaneLayout,
+ int64_t innerLaneData) {
+ SmallVector<int32_t> laneLayout(rank, 1);
+ SmallVector<int32_t> laneData(rank, 1);
+ 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);
+ }
+
+ /// 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();
+ }
+
+ /// True when `op` (a gather load or scatter store) is tied to a
+ /// `vector.multi_reduction`: a load whose result feeds a reduction, or a
+ /// store whose stored value comes from one (through layout-neutral /
+ /// elementwise / insert glue). Coalescing such an access is gated off here;
+ /// it requires reduction-aware layout handling added in follow-up PRs.
+ template <typename OpTy>
+ static bool isReductionTied(OpTy op) {
+ if constexpr (std::is_same_v<OpTy, xegpu::StoreScatterOp>) {
+ SmallVector<Value, 8> worklist{op.getValue()};
+ llvm::SmallPtrSet<Operation *, 16> seen;
+ unsigned steps = 0;
+ while (!worklist.empty() && steps++ < 64) {
+ Value v = worklist.pop_back_val();
+ Operation *def = v.getDefiningOp();
+ if (!def || !seen.insert(def).second)
+ continue;
+ if (isa<vector::MultiDimReductionOp>(def))
+ return true;
+ if (isa<vector::ShapeCastOp, vector::BitCastOp, xegpu::ConvertLayoutOp,
+ vector::InsertOp, vector::InsertStridedSliceOp>(def) ||
+ OpTrait::hasElementwiseMappableTraits(def))
+ for (Value operand : def->getOperands())
+ if (isa<VectorType>(operand.getType()))
+ worklist.push_back(operand);
+ }
+ return false;
+ } else {
+ for (Operation *user : op->getUsers()) {
+ Operation *u = user;
+ while (
+ u &&
+ isa<vector::ShapeCastOp, vector::BitCastOp, xegpu::ConvertLayoutOp>(
+ u)) {
+ if (u->getNumResults() != 1 || u->getResult(0).use_empty())
+ break;
+ u = *u->getResult(0).getUsers().begin();
+ }
+ if (u && isa<vector::MultiDimReductionOp>(u))
+ return true;
+ }
+ return false;
+ }
+ }
+
+ /// Turn a stamped `contiguous_chunk` on `op` into a lane_layout / lane_data /
+ /// inst_data layout, capped by `maxChunkSize`, then remove the attribute.
+ /// Skips ops with a non-uniform mask, an existing lane_data, or a reduction
+ /// tie — these are coalescing concerns, not properties of the offsets.
+ template <typename OpTy>
+ static void applyContiguousChunk(OpTy op, unsigned maxChunkSize) {
+ std::optional<uint64_t> chunk = op.getContiguousChunk();
+ if (!chunk)
+ return;
+ auto cleanup = llvm::scope_exit([&] { op.removeContiguousChunkAttr(); });
+
+ auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
+ auto valueTy = op.getValueType();
+ if (!offsetsTy || !valueTy || offsetsTy.getNumElements() <= 1)
+ return;
+ if (!isAllTrueMask(op.getMask()))
+ return;
+ if (auto layout = op.getLayoutAttr())
+ if (!layout.getEffectiveLaneDataAsInt().empty())
+ return;
+ if (isReductionTied(op))
+ return;
+
+ int64_t inner = offsetsTy.getShape().back();
+ unsigned subgroupSize = lookupSubgroupSize(op);
+ // lane_layout default: min(subgroupSize, inner) rounded to a divisor.
+ int64_t laneLayout =
+ largestPow2Divisor(inner, std::min<int64_t>(subgroupSize, inner));
+ if (laneLayout < 1)
+ return;
+ int64_t perLane = inner / laneLayout;
+ // lane_data = min(contiguity, maxChunkSize, perLane), pow2 divisor of
+ // perLane.
+ int64_t bound =
+ std::min<int64_t>({static_cast<int64_t>(*chunk),
+ static_cast<int64_t>(maxChunkSize), perLane});
+ int64_t factor = largestPow2Divisor(perLane, bound);
+ if (factor < 2)
+ return;
+
+ op.setLayoutAttr(buildLaneDataLayout(op.getContext(), valueTy.getRank(),
+ laneLayout, factor));
}
};
>From b54ebdad62616e3d7ed33173e5321950badaf22c Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Wed, 24 Jun 2026 15:41:43 +0000
Subject: [PATCH 7/7] Fix build/format after rename to XeGPUContiguityAnalysis.
- CMakeLists.txt referenced the old XeGPUCoalesceGatherScatter.cpp name,
breaking the build (CI failure).
- clang-format VectorToXeGPU.cpp (reflow after the longer
contiguous_chunk builder label).
- Update the renamed file's banner comment.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp | 14 +++++++-------
mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt | 2 +-
.../XeGPU/Transforms/XeGPUContiguityAnalysis.cpp | 2 +-
3 files changed, 9 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index a428a9d82e618..bcb2b3530e2f7 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -527,13 +527,13 @@ static LogicalResult lowerToScatteredStoreOp(vector::TransferWriteOp writeOp,
Value mask = vector::ConstantMaskOp::create(
rewriter, loc, VectorType::get(vectorShape, rewriter.getI1Type()),
vectorShape);
- xegpu::StoreScatterOp::create(rewriter, loc, writeOp.getVector(), flatMemref,
- localOffsets, mask,
- /*chunk_size=*/IntegerAttr{},
- /*l1_hint=*/xegpu::CachePolicyAttr{},
- /*l2_hint=*/xegpu::CachePolicyAttr{},
- /*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr, /*contiguous_chunk=*/nullptr);
+ xegpu::StoreScatterOp::create(
+ rewriter, loc, writeOp.getVector(), flatMemref, localOffsets, mask,
+ /*chunk_size=*/IntegerAttr{},
+ /*l1_hint=*/xegpu::CachePolicyAttr{},
+ /*l2_hint=*/xegpu::CachePolicyAttr{},
+ /*l3_hint=*/xegpu::CachePolicyAttr{},
+ /*layout=*/nullptr, /*contiguous_chunk=*/nullptr);
rewriter.eraseOp(writeOp);
return success();
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt b/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
index 1b013400cff34..2ed81ae05ab34 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
@@ -1,7 +1,7 @@
add_mlir_dialect_library(MLIRXeGPUTransforms
XeGPUArrayLengthOptimization.cpp
XeGPUBlocking.cpp
- XeGPUCoalesceGatherScatter.cpp
+ XeGPUContiguityAnalysis.cpp
XeGPUSgToLaneDistribute.cpp
XeGPUUnroll.cpp
XeGPUWgToSgDistribute.cpp
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp
index 1c67503f7445a..5efbed59f38ee 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp
@@ -1,4 +1,4 @@
-//===- XeGPUCoalesceGatherScatter.cpp - Coalesce scatter accesses --------===//
+//===- XeGPUContiguityAnalysis.cpp - Offset contiguity analysis ---------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
More information about the Mlir-commits
mailing list