[llvm] [mlir] [XeGPU][Transform] Add XeGPU array length optimization pass (PR #194062)

[Jianhui Li via llvm-commits]

================
@@ -0,0 +1,252 @@
+//===- XeGPUArrayLengthOptimization.cpp - Array Length Opt -----*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
+#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Debug.h"
+
+namespace mlir {
+namespace xegpu {
+#define GEN_PASS_DEF_XEGPUARRAYLENGTHOPTIMIZATION
+#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"
+} // namespace xegpu
+} // namespace mlir
+
+#define DEBUG_TYPE "xegpu-array-length-optimization"
+
+using namespace mlir;
+
+namespace {
+
+// Subgroup size is typically 16 for Intel GPUs
+constexpr int64_t SUBGROUP_SIZE = 16;
+
+/// Helper to compute array_length from FCD and subgroup size
+static int64_t computeArrayLength(int64_t fcdSize) {
+  if (fcdSize <= SUBGROUP_SIZE)
+    return 1;
+  return fcdSize / SUBGROUP_SIZE;
+}
+
+/// Helper to compute new FCD after introducing array_length
+static int64_t computeNewFCD(int64_t oldFCD, int64_t arrayLength) {
+  return oldFCD / arrayLength;
+}
+
+/// Check if a load_nd or prefetch_nd operation needs optimization
+static bool needsOptimization(xegpu::TensorDescType tdescType) {
+  auto shape = tdescType.getShape();
+  if (shape.size() != 2)
+    return false; // Only 2D tensors
+
+  int64_t fcd = shape[1];
+  if (fcd <= SUBGROUP_SIZE || fcd % SUBGROUP_SIZE != 0)
+    return false; // FCD must be > subgroup_size and evenly divisible
+
+  return tdescType.getArrayLength() == 1; // Skip if already optimized
+}
+
+/// Pattern to rewrite xegpu.create_nd_tdesc operations using simple
+/// RewritePattern
+class OptimizeCreateNdDescOp : public OpRewritePattern<xegpu::CreateNdDescOp> {
+public:
+  using OpRewritePattern<xegpu::CreateNdDescOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(xegpu::CreateNdDescOp op,
+                                PatternRewriter &rewriter) const override {
+    auto tdescType = op.getType();
+    if (!needsOptimization(tdescType))
+      return failure();
+
+    auto shape = tdescType.getShape();
+    int64_t oldFCD = shape[1];
+    int64_t arrayLength = computeArrayLength(oldFCD);
+    int64_t newFCD = computeNewFCD(oldFCD, arrayLength);
----------------
Jianhui-Li wrote:

why create one line function utility with vague name? I have to skip the screen back and force to discover the exact meaning of "new": meaning divide the innermost dim by array length. I dont' think it is worth the utility. And please avoid use "old" and "new" since it doesn't convey much the relation. 

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