[Mlir-commits] [mlir] [MLIR][XeVM] Update HandleVectorExtract pattern. (PR #191052)

[Sang Ik Lee]

https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/191052

>From 5ae6c0bef9f76fcf561c7cf91cc40044f4734d41 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 8 Apr 2026 21:06:23 +0000
Subject: [PATCH 1/2] [MLIR][XeVM] Update HandleVectorExtract pattern. Split
 loads only if pointer address space is private.

---
 mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index e6acc0525fdd5..034692653a387 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -1249,6 +1249,9 @@ class HandleVectorExtractPattern
         // 3. Merge with load as a smaller load
         auto loadOp = cast<LLVM::LoadOp>(srcOp);
         auto loadPtr = loadOp.getAddr();
+        auto loadAddrSpace = loadPtr.getType().getAddressSpace();
+        if (loadAddrSpace != 0)
+          return failure();
         auto loadTy = dyn_cast<VectorType>(loadOp.getType());
         auto elemTy = loadTy.getElementType();
         auto firstIndex = mask[0];
@@ -1257,8 +1260,7 @@ class HandleVectorExtractPattern
         if (firstIndex) {
           auto newPtr = LLVM::GEPOp::create(
               rewriter, loc,
-              LLVM::LLVMPointerType::get(rewriter.getContext(),
-                                         loadPtr.getType().getAddressSpace()),
+              LLVM::LLVMPointerType::get(rewriter.getContext(), loadAddrSpace),
               elemTy, loadPtr, ArrayRef<LLVM::GEPArg>{firstIndex});
           auto newLoad = LLVM::LoadOp::create(rewriter, loc, newVecTy, newPtr);
           rewriter.replaceOp(op, newLoad);

>From 6e742d046fc24f32a70b6cbac4edd2f325b4cfce Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Fri, 10 Apr 2026 01:22:02 +0000
Subject: [PATCH 2/2] Add test case with non private address space.

---
 .../XeVMToLLVM/legalize_large_vector.mlir     | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir b/mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir
index bf11758c074e9..b9721efd6e85e 100644
--- a/mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir
+++ b/mlir/test/Conversion/XeVMToLLVM/legalize_large_vector.mlir
@@ -73,3 +73,23 @@ module @test_match_fail {
     llvm.return
   }
 }
+
+// -----
+
+module @test_non_private_addrspace {
+  // CHECK-LABEL: llvm.func @test_non_private_addrspace
+  // CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>, %[[ARG1:.*]]: !llvm.ptr<1>, %[[ARG2:.*]]: !llvm.ptr<1>
+  llvm.func @test_non_private_addrspace(%arg0: !llvm.ptr<1>, %arg1: !llvm.ptr<1>, %arg2: !llvm.ptr<1>) {
+    // CHECK: %[[VAR0:.*]] = llvm.load %[[ARG0]] : !llvm.ptr<1> -> vector<8xi16>
+    %0 = llvm.load %arg0 : !llvm.ptr<1> -> vector<8xi16>
+    // CHECK: %[[VAR1:.*]] = llvm.shufflevector %[[VAR0]], %[[VAR0]] [0, 1, 2, 3] : vector<8xi16>
+    %1 = llvm.shufflevector %0, %0 [0, 1, 2, 3] : vector<8xi16>
+    // CHECK: %[[VAR2:.*]] = llvm.shufflevector %[[VAR0]], %[[VAR0]] [4, 5, 6, 7] : vector<8xi16>
+    %2 = llvm.shufflevector %0, %0 [4, 5, 6, 7] : vector<8xi16>
+    // CHECK: llvm.store %[[VAR1]], %[[ARG1]] : vector<4xi16>, !llvm.ptr<1>
+    llvm.store %1, %arg1 : vector<4xi16>, !llvm.ptr<1>
+    // CHECK: llvm.store %[[VAR2]], %[[ARG2]] : vector<4xi16>, !llvm.ptr<1>
+    llvm.store %2, %arg2 : vector<4xi16>, !llvm.ptr<1>
+    llvm.return
+  }
+}