[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;
----------------
charithaintc wrote:

why not use `numeric_limits`

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