[llvm] b0cc47c - [InstCombine] Remove scalable vector extracts to and from the same type (#69702)

[via llvm-commits]

Author: Kerry McLaughlin
Date: 2023-10-23T11:21:49+01:00
New Revision: b0cc47c959cb4c4ab54516ee73ede52d208a739f

URL: https://github.com/llvm/llvm-project/commit/b0cc47c959cb4c4ab54516ee73ede52d208a739f
DIFF: https://github.com/llvm/llvm-project/commit/b0cc47c959cb4c4ab54516ee73ede52d208a739f.diff

LOG: [InstCombine] Remove scalable vector extracts to and from the same type (#69702)

visitCallInst already looks for fixed width vector extracts where number of
elements in the source and destination types are equal. This patch modifies
the function to also identify scalable extracts which can be removed.

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
    llvm/test/Transforms/InstCombine/canonicalize-vector-extract.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index e29fb869686ca0b..60233f01e0bde2b 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -2997,24 +2997,27 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
         return replaceOperand(CI, 0, InsertTuple);
     }
 
-    auto *DstTy = dyn_cast<FixedVectorType>(ReturnType);
-    auto *VecTy = dyn_cast<FixedVectorType>(Vec->getType());
+    auto *DstTy = dyn_cast<VectorType>(ReturnType);
+    auto *VecTy = dyn_cast<VectorType>(Vec->getType());
 
-    // Only canonicalize if the destination vector and Vec are fixed
-    // vectors.
     if (DstTy && VecTy) {
-      unsigned DstNumElts = DstTy->getNumElements();
-      unsigned VecNumElts = VecTy->getNumElements();
+      auto DstEltCnt = DstTy->getElementCount();
+      auto VecEltCnt = VecTy->getElementCount();
       unsigned IdxN = cast<ConstantInt>(Idx)->getZExtValue();
 
       // Extracting the entirety of Vec is a nop.
-      if (VecNumElts == DstNumElts) {
+      if (DstEltCnt == VecTy->getElementCount()) {
         replaceInstUsesWith(CI, Vec);
         return eraseInstFromFunction(CI);
       }
 
+      // Only canonicalize to shufflevector if the destination vector and
+      // Vec are fixed vectors.
+      if (VecEltCnt.isScalable() || DstEltCnt.isScalable())
+        break;
+
       SmallVector<int, 8> Mask;
-      for (unsigned i = 0; i != DstNumElts; ++i)
+      for (unsigned i = 0; i != DstEltCnt.getKnownMinValue(); ++i)
         Mask.push_back(IdxN + i);
 
       Value *Shuffle = Builder.CreateShuffleVector(Vec, Mask);

diff  --git a/llvm/test/Transforms/InstCombine/canonicalize-vector-extract.ll b/llvm/test/Transforms/InstCombine/canonicalize-vector-extract.ll
index 26af633698f3552..0b7f1ebc040c2a1 100644
--- a/llvm/test/Transforms/InstCombine/canonicalize-vector-extract.ll
+++ b/llvm/test/Transforms/InstCombine/canonicalize-vector-extract.ll
@@ -10,6 +10,7 @@ declare <3 x i32> @llvm.vector.extract.v3i32.v8i32(<8 x i32> %vec, i64 %idx)
 declare <4 x i32> @llvm.vector.extract.v4i32.nxv4i32(<vscale x 4 x i32> %vec, i64 %idx)
 declare <4 x i32> @llvm.vector.extract.v4i32.v8i32(<8 x i32> %vec, i64 %idx)
 declare <8 x i32> @llvm.vector.extract.v8i32.v8i32(<8 x i32> %vec, i64 %idx)
+declare <vscale x 8 x i32> @llvm.vector.extract.nxv8i32.nxv8i32(<vscale x 8 x i32> %vec, i64 %idx)
 
 ; ============================================================================ ;
 ; Trivial cases
@@ -24,6 +25,15 @@ define <8 x i32> @trivial_nop(<8 x i32> %vec) {
   ret <8 x i32> %1
 }
 
+define <vscale x 8 x i32> @trivial_nop_scalable(<vscale x 8 x i32> %vec) {
+; CHECK-LABEL: define <vscale x 8 x i32> @trivial_nop_scalable(
+; CHECK-SAME: <vscale x 8 x i32> [[VEC:%.*]]) {
+; CHECK-NEXT:    ret <vscale x 8 x i32> [[VEC]]
+;
+  %ext = call <vscale x 8 x i32> @llvm.vector.extract.nxv8i32.nxv8i32(<vscale x 8 x i32> %vec, i64 0)
+  ret <vscale x 8 x i32> %ext
+}
+
 ; ============================================================================ ;
 ; Valid canonicalizations
 ; ============================================================================ ;