[Mlir-commits] [mlir] dd5357d - [mlir][XeGPU][Transform] Add XeGPU contiguity analysis. (#201684)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jun 26 09:15:25 PDT 2026
Author: Md Abdullah Shahneous Bari
Date: 2026-06-26T11:15:20-05:00
New Revision: dd5357d38d6b73e3a687bcc5ea8cb3a858fb3fea
URL: https://github.com/llvm/llvm-project/commit/dd5357d38d6b73e3a687bcc5ea8cb3a858fb3fea
DIFF: https://github.com/llvm/llvm-project/commit/dd5357d38d6b73e3a687bcc5ea8cb3a858fb3fea.diff
LOG: [mlir][XeGPU][Transform] Add XeGPU contiguity analysis. (#201684)
Add an AxisInfo-based (borrows the idea from Triton Axis Info analysis)
dataflow analysis that computes, for each
`xegpu.load` / `xegpu.store` gather/scatter, how many elements are
contiguous
along the innermost offsets dimension, and stamps that count as a
`contiguity` **operation attribute** (`OptionalAttr<I64Attr>`) on the
op.
`contiguity` is a target-independent property of the offsets, not a
request tied to any optimization — a consumer is free to use or ignore
it. The
analysis performs no rewrite. Turning the property into a concrete
`lane_layout` / `lane_data` split (which needs the subgroup size) and
the
actual memory-message rewrite are consumer concerns, handled by later
layout-propagation steps (subsequent PRs) or, for testing, by the apply
helper
in the test pass (`-test-xegpu-coalesce-gather-scatter`).
---------
Co-authored-by: Claude Opus 4.8 <noreply at anthropic.com>
Added:
mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp
mlir/test/Dialect/XeGPU/contiguity-analysis.mlir
mlir/test/Dialect/XeGPU/test-xegpu-coalesce-gather-scatter.mlir
Modified:
mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
mlir/test/Dialect/XeGPU/invalid.mlir
mlir/test/Dialect/XeGPU/ops.mlir
mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index af8c30742bcb2..7f8389a6acc47 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<I64Attr>:$contiguity);
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<I64Attr>:$contiguity);
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 980b4dfab00cd..388bd6145df21 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
@@ -88,6 +88,18 @@ void populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
TypeConverter &typeConverter, RewritePatternSet &patterns,
ConversionTarget &target, Operation *topLevelOp);
+//===----------------------------------------------------------------------===//
+// Coalesce gather/scatter analysis + apply.
+//===----------------------------------------------------------------------===//
+
+/// Run the AxisInfo-based contiguity analysis over `root` and stamp a
+/// `contiguity` 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 `contiguity` attribute are left untouched.
+void runContiguityAnalysis(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/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 9038bf35b6b15..9d92f966c126d 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -499,7 +499,7 @@ static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr);
+ /*layout=*/nullptr, /*contiguity=*/nullptr);
rewriter.replaceOp(readOp, gatherOp.getResult());
return success();
@@ -534,7 +534,7 @@ static LogicalResult lowerToScatteredStoreOp(vector::TransferWriteOp writeOp,
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr);
+ /*layout=*/nullptr, /*contiguity=*/nullptr);
rewriter.eraseOp(writeOp);
return success();
}
@@ -790,7 +790,7 @@ struct GatherLowering : public OpRewritePattern<vector::GatherOp> {
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr);
+ /*layout=*/nullptr, /*contiguity=*/nullptr);
auto selectOp =
arith::SelectOp::create(rewriter, loc, gatherOp.getMask(),
@@ -825,7 +825,8 @@ struct ScatterLowering : public OpRewritePattern<vector::ScatterOp> {
/*l1_hint=*/xegpu::CachePolicyAttr{},
/*l2_hint=*/xegpu::CachePolicyAttr{},
/*l3_hint=*/xegpu::CachePolicyAttr{},
- /*layout=*/nullptr);
+ /*layout=*/nullptr,
+ /*contiguity=*/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..2ffe883eb0d9a 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -113,6 +113,28 @@ isValidGatherScatterBufferParams(Type offsetsTy, Type maskTy,
return success();
}
+// Validates the `contiguity` attribute against the op's offsets type: the
+// innermost offsets dimension is contiguous in runs of `size`, so `size` must
+// be >= 2 and must divide that dimension.
+static LogicalResult
+isValidContiguity(std::optional<uint64_t> contiguity, Type offsetsTy,
+ function_ref<InFlightDiagnostic()> emitError) {
+ if (!contiguity)
+ return success();
+ auto offsetsVecTy = dyn_cast<VectorType>(offsetsTy);
+ if (!offsetsVecTy)
+ return emitError() << "contiguity requires vector offsets (one per lane).";
+ int64_t size = static_cast<int64_t>(*contiguity);
+ int64_t inner = offsetsVecTy.getShape().back();
+ if (size < 2)
+ return emitError() << "contiguity = " << size << " (must be >= 2)";
+ if (inner % size != 0)
+ return emitError() << "contiguity = " << size
+ << " (must divide the innermost offsets dim " << inner
+ << ")";
+ return success();
+}
+
LogicalResult
IsValidMatrixOpParams(VectorType dataTy, MemDescType mdescTy,
UnitAttr subgroup_block_io, DistributeLayoutAttr layout,
@@ -582,6 +604,9 @@ LogicalResult LoadGatherOp::verify() {
}
auto offsetsTy = getOffsets().getType();
+ if (failed(isValidContiguity(getContiguity(), offsetsTy,
+ [&]() { return emitOpError(); })))
+ return failure();
return isValidGatherScatterBufferParams(offsetsTy, maskTy, valueTy, chunkSize,
[&]() { return emitOpError(); });
}
@@ -599,7 +624,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);
+ l2_hint, l3_hint, /*anchor_layout=*/nullptr,
+ /*contiguity=*/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, /*contiguity=*/nullptr);
}
//===----------------------------------------------------------------------===//
@@ -648,6 +674,9 @@ LogicalResult StoreScatterOp::verify() {
}
auto offsetsTy = getOffsets().getType();
+ if (failed(isValidContiguity(getContiguity(), 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, /*contiguity=*/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, /*contiguity=*/nullptr);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt b/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
index 37922f7ef7d24..2ed81ae05ab34 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
+ 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
new file mode 100644
index 0000000000000..3cc8152561681
--- /dev/null
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUContiguityAnalysis.cpp
@@ -0,0 +1,927 @@
+//===- 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.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// 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`.
+//
+// Contiguity is a target-independent property of the offsets.
+//
+// The analysis tracks per-axis information for vectors of integer / index
+// 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"
+#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/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/Support/MathExtras.h"
+#include <numeric>
+#include <optional>
+
+#define DEBUG_TYPE "xegpu-contiguity-analysis"
+
+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` collapses this to the truth. It is not
+/// `numeric_limits<int64_t>::max()` because divisibility values are multiplied
+/// (e.g. in `visitMul`); `1 << 30` keeps those products well within `int64_t`.
+static constexpr int64_t kAxisInfoTop = 1LL << 30;
+
+/// Per-dimension axis information for an SSA vector value of integer / index
+/// type. The fields describe, for each dimension `d`, the pattern of the values
+/// along `d` (examples use a 1-D vector, so `d` is the only / innermost dim):
+/// - `contiguity[d]`: longest run that increases by exactly 1.
+/// `[0, 1, 2, 3]` -> 4; `[0, 1, 0, 1]` -> 2.
+/// - `constancy[d]`: longest run of equal values.
+/// `[5, 5, 5, 5]` -> 4; `[5, 5, 6, 6]` -> 2.
+/// - `divisibility[d]`: a power-of-two divisor of every element.
+/// `[8, 16, 24]` -> 8; `[3, 6, 9]` -> 1.
+/// - `knownConstant`: the value, if the whole vector is one constant.
+/// `dense<7>` -> 7; `[0, 1, 2]` -> nullopt.
+/// - `innerStride`: if set, consecutive values along the innermost dim
diff er
+/// by this constant step. `1` is the contiguous case (`[0,1,2,3]`), `0` the
+/// all-equal case (`[5,5,5,5]`); any other value is a strided progression
+/// (`[0,4,8,12]` -> 4) that `contiguity`/`constancy` can't represent (both
+/// read 1). For a multi-dim vector each inner-dim slice is its own
+/// progression; their bases may
diff er, with the shared inner alignment in
+/// `divisibility[innerDim]`.
+///
+/// Only `contiguity[innerDim]` is consumed when stamping, but all dimensions
+/// are tracked because `vector.transpose` / `vector.shape_cast` permute or move
+/// per-dim info between axes, so an intermediate value's outer dims can become
+/// the inner dim of a later value.
+///
+/// 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 (lattice meet). When a value can arrive from several
+ /// paths (e.g. a block argument, or `arith.select`), only what holds on
+ /// *every* path is safe to assume. So we keep the weaker fact per field:
+ /// `min` of each run length (a run is only guaranteed as long as the shortest
+ /// incoming one), `gcd` of divisibility, and a value/stride only when both
+ /// sides agree. This is what makes the contiguity we later stamp sound rather
+ /// than "undecidable" — it is the largest run guaranteed on all paths.
+ 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();
+ }
+
+ // arith.constant. The four cases below, by example:
+ // - scalar int `arith.constant 8 : index`
+ // - non-int scalar `arith.constant 1.0 : f32` (pessimistic)
+ // - splat vector `arith.constant dense<5> : vector<16xindex>`
+ // - dense vector `arith.constant dense<[0,1,2,3]> : vector<4xindex>`
+ LogicalResult visitConstant(arith::ConstantOp op,
+ ArrayRef<AxisInfoLattice *> results) {
+ auto vt = dyn_cast<VectorType>(op.getType());
+ if (!vt) {
+ // Scalar integer, e.g. `arith.constant 8 : index`: a single known value,
+ // contiguity/constancy 1, divisibility from the value (8 -> 8).
+ 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();
+ }
+ // Non-integer scalar, e.g. `arith.constant 1.0 : f32`: nothing to track.
+ setAllPessimistic(op, results);
+ return success();
+ }
+ auto dense = dyn_cast<DenseIntElementsAttr>(op.getValue());
+ if (!dense) {
+ setAllPessimistic(op, results);
+ return success();
+ }
+ auto shape = vt.getShape();
+ // Splat, e.g. `arith.constant dense<5> : vector<16xindex>`: every element
+ // equal, so constancy = full extent, innerStride 0.
+ if (dense.isSplat()) {
+ int64_t c = dense.getSplatValue<APInt>().getSExtValue();
+ AxisInfo v = splatAxisInfo(shape, c);
+ propagateIfChanged(results[0], results[0]->join(v));
+ return success();
+ }
+
+ // General dense vector, e.g. `arith.constant dense<[0,1,2,3]> :
+ // vector<4xindex>`. Compute innermost-dim contiguity / constancy /
+ // base-divisibility by iterating the dense values along the inner stride
+ // (here stride 1 -> contiguity 4). 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
diff erent 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 extent) {
+ return a.knownConstant && *a.knownConstant == 1 &&
+ a.constancy[d] >= extent;
+ };
+ 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`. 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`.
+ //
+ // 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).
+ //
+ // 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) {
+ 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
+ // 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) {
+ 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;
+
+//===----------------------------------------------------------------------===//
+// Analysis driver.
+//===----------------------------------------------------------------------===//
+
+/// Stamp a `contiguity` 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 void analyzeAndStampContiguity(OpTy op, DataFlowSolver &solver) {
+ auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
+ if (!offsetsTy || offsetsTy.getNumElements() <= 1)
+ return;
+ // A pre-existing `contiguity` (user-authored, or stamped by an earlier run)
+ // takes precedence; leave it untouched so the analysis is idempotent.
+ if (op.getContiguity())
+ return;
+ const auto *lat = solver.lookupState<AxisInfoLattice>(op.getOffsets());
+ if (!lat || !lat->getValue().isInitialized())
+ return;
+ const AxisInfo &info = lat->getValue();
+ unsigned innerDim = offsetsTy.getRank() - 1;
+ int64_t inner = offsetsTy.getShape()[innerDim];
+ // The attribute records a contiguity that tiles the inner dim, so it must
+ // divide it (verified on the op). Round the measured run length down to the
+ // largest divisor of `inner` that does not exceed it.
+ int64_t contiguity = std::min<int64_t>(info.contiguity[innerDim], inner);
+ while (contiguity >= 2 && inner % contiguity != 0)
+ --contiguity;
+ if (contiguity < 2)
+ return;
+ op.setContiguity(contiguity);
+}
+
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Public API.
+//===----------------------------------------------------------------------===//
+
+void mlir::xegpu::runContiguityAnalysis(Operation *root) {
+ DataFlowSolver solver;
+ solver.load<dataflow::DeadCodeAnalysis>();
+ solver.load<mlir::xegpu::detail::axis_dataflow::AxisInfoAnalysis>();
+ 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 an attribute.
+ root->walk([&](Operation *op) {
+ if (auto load = dyn_cast<xegpu::LoadGatherOp>(op))
+ analyzeAndStampContiguity(load, solver);
+ else if (auto store = dyn_cast<xegpu::StoreScatterOp>(op))
+ analyzeAndStampContiguity(store, solver);
+ });
+}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 23d3877ea63dd..f9a82bf2b1684 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -505,7 +505,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, /*contiguity=*/nullptr);
Value result = newOp->getResult(0);
if (distResultTy1D != distResultTy)
@@ -1037,7 +1037,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,
+ /*contiguity=*/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 a9c73b3b84025..74c358cef90df 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -765,7 +765,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,
+ /*contiguity=*/nullptr);
newOps.push_back(newOp);
}
@@ -860,7 +861,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,
+ /*contiguity=*/nullptr);
}
rewriter.eraseOp(op);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 05220833b07d7..9ca6b3c2b0272 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -832,8 +832,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);
+ op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(), newLayout,
+ /*contiguity=*/nullptr);
newLoadOps.push_back(newLoadOp);
}
rewriter.replaceOpWithMultiple(op, {newLoadOps});
@@ -879,7 +879,8 @@ struct WgToSgStoreScatterOp
xegpu::StoreScatterOp::create(rewriter, loc, val, op.getDest(), offs,
mask, chunkSizeAttr, op.getL1HintAttr(),
op.getL2HintAttr(), op.getL3HintAttr(),
- layout.dropSgLayoutAndData());
+ layout.dropSgLayoutAndData(),
+ /*contiguity=*/nullptr);
}
rewriter.eraseOp(op);
return success();
diff --git a/mlir/test/Dialect/XeGPU/contiguity-analysis.mlir b/mlir/test/Dialect/XeGPU/contiguity-analysis.mlir
new file mode 100644
index 0000000000000..f223b325c2f9e
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/contiguity-analysis.mlir
@@ -0,0 +1,256 @@
+// RUN: mlir-opt -split-input-file \
+// RUN: -test-xegpu-coalesce-gather-scatter="analyze-only=true" %s | FileCheck %s
+
+// Contiguity analysis: stamps the `contiguity` attribute on gather/scatter ops
+// whose offsets are contiguous (runs of >= 2) along the innermost dimension.
+// The stamped value is the inner-dim contiguity, rounded down to a divisor of
+// the inner extent. The analysis is target-independent and mask-independent;
+// the lane_data split and any access-pattern gating live in the consumer.
+
+// -----
+// 1-D vector.step -> stride-1, fully contiguous over the 32-element inner dim.
+// CHECK-LABEL: func.func @load_step_offsets(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 32 : i64}>
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+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
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// Dense stride-1 constant offsets -> contiguity 32.
+// CHECK-LABEL: func.func @load_dense_ap_offsets(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 32 : i64}>
+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
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xi32>
+ return %v : vector<32xi32>
+}
+
+// -----
+// Stride-4 offsets: not contiguous, no attribute stamped.
+// CHECK-LABEL: func.func @load_stride4_no_attr(
+// CHECK: xegpu.load
+// CHECK-NOT: contiguity
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+func.func @load_stride4_no_attr(%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
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// All-equal offsets: inner dim is constant, not contiguous -> no attribute.
+// CHECK-LABEL: func.func @load_broadcast_offsets_no_attr(
+// CHECK: xegpu.load
+// CHECK-NOT: contiguity
+func.func @load_broadcast_offsets_no_attr(%ptr: i64) -> vector<32xf32> {
+ %offsets = arith.constant dense<0> : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// The analysis is mask-independent: a partial mask still gets the attribute
+// (the mask check is a consumer concern).
+// CHECK-LABEL: func.func @load_partial_mask(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 32 : i64}>
+func.func @load_partial_mask(%ptr: i64, %mask: vector<32xi1>) -> vector<32xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %v = xegpu.load %ptr[%offsets], %mask
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// Store with vector.step offsets -> contiguity 32 on the store.
+// CHECK-LABEL: func.func @store_step_offsets(
+// CHECK: xegpu.store
+// CHECK-SAME: <{contiguity = 32 : i64}>
+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
+ : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
+ return
+}
+
+// -----
+// 2-D leading-unit dim: contiguity measured on the inner dim -> 32.
+// CHECK-LABEL: func.func @load_2d_leading_unit(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 32 : i64}>
+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
+ : i64, vector<1x32xindex>, vector<1x32xi1> -> vector<1x32xf32>
+ return %v : vector<1x32xf32>
+}
+
+// -----
+// True 2-D dense AP: each row stride-1 over 16 -> contiguity 16.
+// CHECK-LABEL: func.func @load_2d_dense_ap(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 16 : i64}>
+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
+ : i64, vector<2x16xindex>, vector<2x16xi1> -> vector<2x16xf32>
+ return %v : vector<2x16xf32>
+}
+
+// -----
+// 2-D dense values whose inner row is not a stride-1 AP: no attribute.
+// CHECK-LABEL: func.func @load_2d_non_ap(
+// CHECK: xegpu.load
+// CHECK-NOT: contiguity
+func.func @load_2d_non_ap(%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
+ : i64, vector<2x16xindex>, vector<2x16xi1> -> vector<2x16xf32>
+ return %v : vector<2x16xf32>
+}
+
+// -----
+// `divui` by a constant equal to the inner stride recovers stride-1
+// contiguity: (step * 2) / 2 = [0,1,..,31] -> contiguity 32. The dividend is
+// derived from vector.step (not a constant) so the divui is genuinely
+// exercised even if the solver later folds all-constant arith ops.
+// CHECK-LABEL: func.func @load_divui_recovers(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 32 : i64}>
+func.func @load_divui_recovers(%ptr: i64) -> vector<32xf32> {
+ %step = vector.step : vector<32xindex>
+ %c2 = arith.constant dense<2> : vector<32xindex>
+ %even = arith.muli %step, %c2 : vector<32xindex>
+ %offsets = arith.divui %even, %c2 : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// `divui` by a constant that does not divide the inner stride: not recovered.
+// CHECK-LABEL: func.func @load_divui_non_divisor(
+// CHECK: xegpu.load
+// CHECK-NOT: contiguity
+func.func @load_divui_non_divisor(%ptr: i64) -> vector<16xf32> {
+ %step = vector.step : vector<16xindex>
+ %c2 = arith.constant dense<2> : vector<16xindex>
+ %even = arith.muli %step, %c2 : 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
+ : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ return %v : vector<16xf32>
+}
+
+// -----
+// `remui` collapsing each element to a constant residue: inner dim is uniform,
+// not contiguous -> no attribute.
+// CHECK-LABEL: func.func @load_remui_inner_uniform(
+// CHECK: xegpu.load
+// CHECK-NOT: contiguity
+func.func @load_remui_inner_uniform(%ptr: i64) -> vector<16xf32> {
+ %step = vector.step : vector<16xindex>
+ %c2 = arith.constant dense<2> : vector<16xindex>
+ %even = arith.muli %step, %c2 : vector<16xindex>
+ %offsets = arith.remui %even, %c2 : vector<16xindex>
+ %mask = arith.constant dense<true> : vector<16xi1>
+ %v = xegpu.load %ptr[%offsets], %mask
+ : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ return %v : vector<16xf32>
+}
+
+// -----
+// `shli` then `shrui` cancel: (step << 1) >> 1 -> stride 1 -> contiguity 32.
+// CHECK-LABEL: func.func @load_shli_then_shrui(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 32 : i64}>
+func.func @load_shli_then_shrui(%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
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// `shli` alone scales the stride to 2: not contiguous -> no attribute.
+// CHECK-LABEL: func.func @load_shli_unchanged(
+// CHECK: xegpu.load
+// CHECK-NOT: contiguity
+func.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
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// `select` between two AP arms with matching inner-dim properties preserves
+// contiguity 32.
+// CHECK-LABEL: func.func @load_select_two_aps(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 32 : i64}>
+func.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
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
+
+// -----
+// A pre-existing `contiguity` takes precedence: the analysis leaves it alone.
+// CHECK-LABEL: func.func @user_attr_preserved(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 2 : i64}>
+func.func @user_attr_preserved(%ptr: i64) -> vector<32xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask <{contiguity = 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 4b811873cba81..d5d4950fe7d7e 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -789,3 +789,19 @@ 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 @contiguity_too_small(%src: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
+ // expected-error at +1 {{contiguity = 1 (must be >= 2)}}
+ %val = xegpu.load %src[%offset], %mask <{contiguity = 1 : i64}>
+ : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ return
+}
+
+// -----
+func.func @contiguity_does_not_divide(%src: i64, %offset: vector<6xindex>, %mask: vector<6xi1>) {
+ // expected-error at +1 {{contiguity = 4 (must divide the innermost offsets dim 6)}}
+ %val = xegpu.load %src[%offset], %mask <{contiguity = 4 : 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 0b1ac71cdbd32..6cffa3eec369b 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_contiguity(%[[arg0:.*]]: i64, %[[arg1:.*]]: vector<16xindex>, %[[arg2:.*]]: vector<16xi1>) {
+gpu.func @load_contiguity(%src: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
+ // A user-provided `contiguity` round-trips through the optional op attribute.
+ // CHECK: xegpu.load %[[arg0]][%[[arg1]]], %[[arg2]] <{contiguity = 4 : i64}> : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ %val = xegpu.load %src[%offset], %mask <{contiguity = 4 : i64}>
+ : i64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
+ gpu.return
+}
+
+// CHECK: gpu.func @store_contiguity(%[[arg0:.*]]: vector<16xf32>, %[[arg1:.*]]: i64, %[[arg2:.*]]: vector<16xindex>, %[[arg3:.*]]: vector<16xi1>) {
+gpu.func @store_contiguity(%val: vector<16xf32>, %dest: i64, %offset: vector<16xindex>, %mask: vector<16xi1>) {
+ // CHECK: xegpu.store %[[arg0]], %[[arg1]][%[[arg2]]], %[[arg3]] <{contiguity = 4 : i64}> : vector<16xf32>, i64, vector<16xindex>, vector<16xi1>
+ xegpu.store %val, %dest[%offset], %mask <{contiguity = 4 : i64}>
+ : 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>
diff --git a/mlir/test/Dialect/XeGPU/test-xegpu-coalesce-gather-scatter.mlir b/mlir/test/Dialect/XeGPU/test-xegpu-coalesce-gather-scatter.mlir
new file mode 100644
index 0000000000000..fb32e0012679b
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/test-xegpu-coalesce-gather-scatter.mlir
@@ -0,0 +1,110 @@
+// 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
+
+// Driver test: turns the `contiguity` attribute the analysis stamps into a
+// lane_layout / lane_data layout. The contiguity analysis itself is covered by
+// contiguity-analysis.mlir; here we only check that the driver derives the
+// layout from a stamped (or already-contiguous) access. The subgroup size
+// comes from the target (pvc -> 16), so each case carries an xevm.target.
+
+// -----
+// Contiguous 32-element load, subgroup_size = 16: lane_layout = 16,
+// perLane = 2, lane_data = min(contiguity, max-chunk-size, perLane) = 2.
+// CHECK-LABEL: gpu.func @load_step(
+// 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 [#xevm.target<chip = "pvc">] {
+ gpu.func @load_step(%ptr: i64) -> vector<32xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ gpu.return %v : vector<32xf32>
+ }
+}
+
+// -----
+// Store mirrors the load.
+// CHECK-LABEL: gpu.func @store_step(
+// 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>
+gpu.module @kernel_store [#xevm.target<chip = "pvc">] {
+ gpu.func @store_step(%ptr: i64, %v: vector<32xf32>) {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ xegpu.store %v, %ptr[%offsets], %mask
+ : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
+ gpu.return
+ }
+}
+
+// -----
+// Non-contiguous (stride-4) offsets: the analysis stamps nothing, so the
+// driver leaves the op unchanged.
+// CHECK-LABEL: gpu.func @load_stride4_unchanged(
+// CHECK: xegpu.load
+// CHECK-NOT: lane_data
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+gpu.module @kernel_stride4 [#xevm.target<chip = "pvc">] {
+ gpu.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
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ gpu.return %v : vector<32xf32>
+ }
+}
+
+// -----
+// 2-D leading-unit offsets: lane_layout / lane_data are placed on the inner
+// dim -> [1, 16] / [1, 2].
+// CHECK-LABEL: gpu.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>
+gpu.module @kernel_2d [#xevm.target<chip = "pvc">] {
+ gpu.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
+ : i64, vector<1x32xindex>, vector<1x32xi1> -> vector<1x32xf32>
+ gpu.return %v : vector<1x32xf32>
+ }
+}
+
+// -----
+// max-chunk-size = 4 saturates at perLane = 2 anyway, so lane_data stays 2.
+// CHECK4-LABEL: gpu.func @load_step_chunk4(
+// CHECK4: xegpu.load
+// CHECK4-SAME: layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>
+gpu.module @kernel_chunk4 [#xevm.target<chip = "pvc">] {
+ gpu.func @load_step_chunk4(%ptr: i64) -> vector<32xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ gpu.return %v : vector<32xf32>
+ }
+}
+
+// -----
+// No target: the driver cannot size the lanes, so it leaves the op unchanged
+// (the stamped contiguity is just removed).
+// CHECK-LABEL: func.func @load_no_target_unchanged(
+// CHECK: xegpu.load
+// CHECK-NOT: lane_data
+// CHECK-NOT: contiguity
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+func.func @load_no_target_unchanged(%ptr: i64) -> vector<32xf32> {
+ %offsets = vector.step : vector<32xindex>
+ %mask = arith.constant dense<true> : vector<32xi1>
+ %v = xegpu.load %ptr[%offsets], %mask
+ : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ return %v : vector<32xf32>
+}
diff --git a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
index 6ea7d82295cf9..f6b0c50da91dd 100644
--- a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
+++ b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
@@ -16,6 +16,7 @@
#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"
@@ -446,6 +447,105 @@ 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 driver that turns the contiguity attribute into a lane_data "
+ "layout on gather/scatter ops.";
+ }
+
+ 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 the contiguity attribute); "
+ "do not apply."),
+ llvm::cl::init(false)};
+
+ void runOnOperation() override {
+ xegpu::runContiguityAnalysis(getOperation());
+ if (analyzeOnly)
+ return;
+ getOperation()->walk([&](Operation *op) {
+ if (auto load = dyn_cast<xegpu::LoadGatherOp>(op))
+ applyContiguity(load, maxChunkSize);
+ else if (auto store = dyn_cast<xegpu::StoreScatterOp>(op))
+ applyContiguity(store, maxChunkSize);
+ });
+ }
+
+private:
+ /// 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);
+ }
+
+ /// Minimal driver: read the `contiguity` attribute the analysis stamped and
+ /// turn it into a `lane_data` layout. This is only a stand-in for the real
+ /// consumer (layout propagation) so the analysis output can be checked
+ /// end-to-end; it handles just the simple power-of-two case.
+ template <typename OpTy>
+ static void applyContiguity(OpTy op, unsigned maxChunkSize) {
+ std::optional<uint64_t> contiguity = op.getContiguity();
+ if (!contiguity)
+ return;
+ op.removeContiguityAttr();
+
+ auto offsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
+ auto valueTy = op.getValueType();
+ if (!offsetsTy || !valueTy)
+ return;
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
+ if (!uArch)
+ return;
+ int64_t subgroupSize = uArch->getSubgroupSize();
+
+ // Tile size and subgroup size are powers of two, so the smaller is already
+ // the largest power-of-two divisor; min() suffices throughout.
+ int64_t inner = offsetsTy.getShape().back();
+ int64_t laneLayout = std::min<int64_t>(subgroupSize, inner);
+ int64_t perLane = inner / laneLayout;
+ int64_t laneData =
+ std::min<int64_t>({static_cast<int64_t>(*contiguity),
+ static_cast<int64_t>(maxChunkSize), perLane});
+ if (laneData < 2)
+ return;
+
+ op.setLayoutAttr(buildLaneDataLayout(op.getContext(), valueTy.getRank(),
+ laneLayout, laneData));
+ }
+};
+
} // namespace
namespace mlir {
@@ -458,6 +558,7 @@ void registerTestXeGPULowerings() {
PassRegistration<TestXeGPUPropagateLayouts>();
PassRegistration<TestXeGPUResolveLayoutConflicts>();
PassRegistration<TestXeGPUArrayLengthOptimization>();
+ PassRegistration<TestXeGPUCoalesceGatherScatter>();
}
} // namespace test
} // namespace mlir
More information about the Mlir-commits
mailing list