[MLIR][XeVM] Update HandleVectorExtractPattern

[silee2]

isExtractContiguousSlice:

• Check if mask size is not greater than the vector size of the operand.
• Check if mask values do not exceed vector size.

HandleVectorExtractPattern:

• Narrow the scope of matching to,
  • Source shuffle doing contiguous extract
  • Source shuffle with at least the same mask size.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Sang Ik Lee (silee2)

Changes

isExtractContiguousSlice:

• Check if mask size is not greater than the vector size of the operand.
• Check if mask values do not exceed vector size.

HandleVectorExtractPattern:

• Narrow the scope of matching to,
  • Source shuffle doing contiguous extract
  • Source shuffle with at least the same mask size.

---

Full diff: https://github.com/llvm/llvm-project/pull/186247.diff

1 Files Affected:

• (modified) mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp (+17-2)

diff --git a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
index 24009b63e8e26..baf9a6ccc4b09 100644
--- a/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
+++ b/mlir/lib/Conversion/XeVMToLLVM/XeVMToLLVM.cpp
@@ -895,11 +895,17 @@ static bool isExtractingContiguousSlice(LLVM::ShuffleVectorOp op) {
   if (op.getV1() != op.getV2())
     return false;
   auto maskAttr = op.getMask();
+  int64_t maskSize = static_cast<int64_t>(maskAttr.size());
+  int64_t sourceSize = op.getV1().getType().getNumElements();
+  if (maskSize > sourceSize)
+    return false;
   int64_t firstIndex = maskAttr[0];
-  for (int64_t i = 1; i < static_cast<int64_t>(maskAttr.size()); ++i) {
+  for (int64_t i = 1; i < maskSize; ++i) {
     int64_t index = maskAttr[i];
     if (index != firstIndex + i)
       return false;
+    if (index >= sourceSize)
+      return false;
   }
   return true;
 }
@@ -984,9 +990,18 @@ class HandleVectorExtractPattern
             LLVM::BitcastOp::create(rewriter, loc, ty, newShuffle);
         rewriter.replaceOp(op, newBitcast);
       } else if (isa<LLVM::ShuffleVectorOp>(srcOp)) {
-        // 2. Merge with another shuffle vector op
+        // 2. Merge with source shuffle vector op if,
+        //  - the source op is also extracting a contigous slice.
+        //  - the source op mask size is not smaller than the current
+        //    op mask size.
+        // And create a new shuffle vector op directly from the source
+        // of the first shuffle.
         auto srcShuffle = cast<LLVM::ShuffleVectorOp>(srcOp);
+        if (!isExtractingContiguousSlice(srcShuffle))
+          return failure();
         auto srcMask = srcShuffle.getMask();
+        if (srcMask.size() < mask.size())
+          return failure();
         SmallVector<int32_t> combinedMask;
         for (auto index : mask) {
           combinedMask.push_back(srcMask[index]);

[Garra1980]

any new or adjusted old tests for the change?

[silee2]

any new or adjusted old tests for the change?

Added a new test where match should fail.

[Jianhui-Li]

LGTM. Thanks!

[Jianhui-Li]

nit: consider using destSize instead of maskSize.