[Mlir-commits] [mlir] [MLIR][XeGPU] Disable block count usage in layout propagation (PR #168504)

[Artem Kroviakov]

https://github.com/akroviakov updated https://github.com/llvm/llvm-project/pull/168504

>From 3358a59c590cb712cfe46edf1e62ab7dcbcee6b0 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Tue, 18 Nov 2025 08:43:56 +0000
Subject: [PATCH 1/2] [MLIR][XeGPU] Disable block count usage in layout
 propagation

---
 mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index b3a780abd3f12..6b3ba5a5981ce 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -495,8 +495,7 @@ void LayoutInfoPropagation::visitPrefetchNdOp(
   auto [bWidth, bHeight, bCount] = blockWHC.value();
   SmallVector<int> instData;
   int instWidth = xegpu::getLargestDivisor(
-      static_cast<int>(tdescTy.getDimSize(tdescTy.getRank() - 1)), bWidth,
-      bCount);
+      static_cast<int>(tdescTy.getDimSize(tdescTy.getRank() - 1)), bWidth);
   if (instWidth == -1)
     prefetch.emitWarning(
         "No suitable instruction multiple found for the given shape.");
@@ -702,8 +701,7 @@ void LayoutInfoPropagation::visitStoreNdOp(
   auto [bWidth, bHeight, bCount] = blockWHC.value();
   SmallVector<int> instData;
   int instWidth = xegpu::getLargestDivisor(
-      static_cast<int>(dataTy.getDimSize(dataTy.getRank() - 1)), bWidth,
-      bCount);
+      static_cast<int>(dataTy.getDimSize(dataTy.getRank() - 1)), bWidth);
   if (instWidth == -1)
     store.emitWarning(
         "No suitable instruction multiple found for the given shape.");

>From 161303c8562796c03c5a720477590fab2d42400a Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Thu, 20 Nov 2025 13:25:45 +0000
Subject: [PATCH 2/2] Add inst option description

---
 mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
index 12270af870b3b..1fb01c3586e93 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
@@ -37,6 +37,18 @@ def XeGPUPropagateLayout : Pass<"xegpu-propagate-layout"> {
     propagate the layouts required for their operands to the producers. With
     this propagated layout information, pass will then update op result type
     with the layout information.
+
+    `inst`:
+      Sets the `inst_data` field of the layout attribute.
+      The goal is to select the highest granularity of the
+      instruction shape to minimize the number of instructions required
+      for processing the user shape.
+      For nd operations, the granularity depends on (W)idth, (H)eight and (B)lock count.
+      It is possible that the max. block count of one operation is different from the
+      max. allowed block count of another operation, whereas both operations
+      may use the same tensor descriptor. The propagation is not the right place for creating a
+      second tdesc with the correct block count and other related ops. Hence,
+      the subsequent optimization passes may decide on if or how to handle the block count.
   }];
   let dependentDialects = ["memref::MemRefDialect", "xegpu::XeGPUDialect",
                            "vector::VectorDialect"];