[Mlir-commits] [mlir] [MLIR][XeGPU] Allow CRI as supported uArch. (PR #197229)

[Sang Ik Lee]

https://github.com/silee2 created https://github.com/llvm/llvm-project/pull/197229

Some XeGPU passes are not aware of CRI uArch and does not handle correctly. Logic for handling is the same as PVC and BMG.
Updated passes are.
Vector to XeGPU
XeGPU to XeVM
XeGPU PeepHole optimizer
XeGPU Utils

>From 7afca0d2cfd4aeeea8cd76f033775050b1909276 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Tue, 12 May 2026 15:15:27 +0000
Subject: [PATCH] [MLIR][XeGPU] Allow CRI as supported uArch. Some XeGPU passes
 are not aware of CRI uArch and does not handle correctly. Logic for handling
 is the same as PVC and BMG. Updated passes are. Vector to XeGPU XeGPU to XeVM
 XeGPU PeepHole optimizer XeGPU Utils

---
 .../lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp |  4 ++--
 mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp    |  7 ++++---
 .../XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp    | 14 +++++++++-----
 mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp        |  3 ++-
 4 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 36e1975017881..d8eca95cbf23c 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -590,7 +590,7 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
     auto chip = xegpu::getChipStr(readOp);
     // Lower to scattered load Op if the target HW doesn't have 2d block load
     // support and the load is not from shared memory.
-    if ((chip != "pvc" && chip != "bmg") ||
+    if ((chip != "pvc" && chip != "bmg" && chip != "cri") ||
         readOp.getVectorType().getRank() > 2) {
 
       // TODO: add support for OutOfBound access
@@ -723,7 +723,7 @@ struct TransferWriteLowering
     auto chip = xegpu::getChipStr(writeOp);
     // Lower to scattered store Op if the target HW doesn't have 2d block
     // store support and the memref is not SLM.
-    if ((chip != "pvc" && chip != "bmg") ||
+    if ((chip != "pvc" && chip != "bmg" && chip != "cri") ||
         writeOp.getVectorType().getRank() > 2) {
 
       // TODO: add support for OutOfBound access
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index d9ff452fb4db3..7e58b19b25660 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -763,10 +763,11 @@ class LoadStoreMatrixToXeVMPattern : public OpConversionPattern<OpType> {
 
     if (valOrResVecTy.getNumElements() >= 1) {
       auto chipOpt = xegpu::getChipStr(op);
-      if (!chipOpt || (*chipOpt != "pvc" && *chipOpt != "bmg")) {
-        // the lowering for chunk load only works for pvc and bmg
+      if (!chipOpt ||
+          (*chipOpt != "pvc" && *chipOpt != "bmg" && *chipOpt != "cri")) {
+        // the lowering for chunk load only works for pvc, bmg or cri
         return rewriter.notifyMatchFailure(
-            op, "The lowering is specific to pvc or bmg.");
+            op, "The lowering is specific to pvc, bmg or cri.");
       }
     }
 
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index 0be3b86120c86..efe8e23978e91 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -256,9 +256,11 @@ class XeGPUCreateNdDescOpPattern final
     // Get the target uArch info.
     auto chipStr = xegpu::getChipStr(createNdOp);
     // Check if the chip is supported.
-    assert(
-        chipStr && (chipStr.value() == "pvc" || chipStr.value() == "bmg") &&
-        "Expecting target chip to be pvc or bmg for transpose optimization.");
+    assert(chipStr &&
+           (chipStr.value() == "pvc" || chipStr.value() == "bmg" ||
+            chipStr.value() == "cri") &&
+           "Expecting target chip to be pvc, bmg or cri for transpose "
+           "optimization.");
     const uArch *targetuArch = xegpu::uArch::getUArch(chipStr.value());
 
     auto convertType = tryOptimize(tdescTy, targetuArch);
@@ -572,12 +574,14 @@ struct XeGPUPeepHoleOptimizerPass final
     bool isTargetSupported = false;
     getOperation()->walk([&](gpu::GPUFuncOp funcOp) {
       auto chipStr = xegpu::getChipStr(funcOp);
-      if (chipStr && (chipStr.value() == "pvc" || chipStr.value() == "bmg"))
+      if (chipStr && (chipStr.value() == "pvc" || chipStr.value() == "bmg" ||
+                      chipStr.value() == "cri"))
         isTargetSupported = true;
     });
 
     if (!isTargetSupported) {
-      DBGS() << "XeGPUPeepHoleOptimizerPass only supports PVC and BMG targets."
+      DBGS() << "XeGPUPeepHoleOptimizerPass only supports PVC, BMG and CRI "
+                "targets."
              << "\n";
       return;
     }
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index d7cbec3b3f85d..71bcf3d53f9b7 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -922,7 +922,8 @@ bool xegpu::requireTranspose(const xegpu::DistributeLayoutAttr layout,
   // Return false for unsupported targets.
   // TODO: Add more support or move to target info.
   if (uArch->getName().equals_insensitive("pvc") &&
-      uArch->getName().equals_insensitive("bmg"))
+      uArch->getName().equals_insensitive("bmg") &&
+      uArch->getName().equals_insensitive("cri"))
     return false;
   if (!layout)
     return false;