[Mlir-commits] [mlir] [mlir][XeGPU][Transform] Add gather/scatter coalescing analysis. (PR #201684)

[Charitha Saumya]

================
@@ -0,0 +1,891 @@
+//===- 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-coalesce-gather-scatter"
+
+using namespace mlir;
+
+// AxisInfo and AxisInfoAnalysis are intentionally placed in a named namespace
+// (not anonymous) so the `dataflow::Lattice<AxisInfo>` template instantiation
+// gets a stable, externally-visible name. The TypeID machinery requires that
+// for `MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID` to apply.
+namespace mlir::xegpu::detail::axis_dataflow {
+
+//===----------------------------------------------------------------------===//
+// AxisInfo: per-axis contiguity / constancy / divisibility lattice.
+//===----------------------------------------------------------------------===//
+
+/// Sentinel "very large" value for unconstrained dimensions. Any real shape
+/// is far smaller, so component-wise `min` will collapse this to the truth.
+static constexpr int64_t kAxisInfoTop = 1LL << 30;
+
+/// Per-dimension axis information for an SSA value of integer / index type.
+/// All fields describe runs of consecutive elements along a dimension `d`:
+///   - `contiguity[d]`: longest run that increases by exactly 1.
+///   - `constancy[d]`: longest run of equal values.
+///   - `divisibility[d]`: a power-of-two divisor of every element.
+///   - `knownConstant`: the value, if the whole vector is one constant.
+///   - `innerStride`: if set, each row along the innermost dim is an
+///     arithmetic progression with this step. `1` is the contiguous case,
+///     `0` the all-equal case; any other value is a strided progression that
+///     `contiguity`/`constancy` can't represent (both read 1). Per-row base
+///     may vary; per-row alignment lives in `divisibility[innerDim]`.
+///
+/// Pessimistic / entry value: contiguity=1, constancy=1, divisibility=1,
+/// innerStride absent.
+struct AxisInfo {
+  SmallVector<int64_t> contiguity;
+  SmallVector<int64_t> constancy;
+  SmallVector<int64_t> divisibility;
+  std::optional<int64_t> knownConstant;
+  std::optional<int64_t> innerStride;
+
+  AxisInfo() = default;
+
+  static AxisInfo getPessimistic(unsigned rank) {
+    AxisInfo v;
+    v.contiguity.assign(rank, 1);
+    v.constancy.assign(rank, 1);
+    v.divisibility.assign(rank, 1);
+    return v;
+  }
+
+  unsigned getRank() const { return contiguity.size(); }
+  bool isInitialized() const { return getRank() > 0; }
+
+  bool operator==(const AxisInfo &rhs) const {
+    return contiguity == rhs.contiguity && constancy == rhs.constancy &&
+           divisibility == rhs.divisibility &&
+           knownConstant == rhs.knownConstant && innerStride == rhs.innerStride;
+  }
+
+  /// Conservative join. Two values reaching the same SSA value via different
+  /// control-flow paths must agree on what holds.
+  static AxisInfo join(const AxisInfo &lhs, const AxisInfo &rhs) {
+    if (!lhs.isInitialized())
+      return rhs;
+    if (!rhs.isInitialized())
+      return lhs;
+    assert(lhs.getRank() == rhs.getRank());
+    AxisInfo out;
+    unsigned r = lhs.getRank();
+    out.contiguity.resize(r);
+    out.constancy.resize(r);
+    out.divisibility.resize(r);
+    for (unsigned d = 0; d < r; ++d) {
+      out.contiguity[d] = std::min(lhs.contiguity[d], rhs.contiguity[d]);
----------------
charithaintc wrote:

huh that makes sense. run means number of consecutive elements.

https://github.com/llvm/llvm-project/pull/201684