| //===- 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 differ |
| /// 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 differ, 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 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 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); |
| }); |
| } |
