[InstCombine] Remove scalable vector extracts to and from the same type

[kmclaughlin-arm]

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.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Kerry McLaughlin (kmclaughlin-arm)

Changes

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.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp (+10-8)
• (modified) llvm/test/Transforms/InstCombine/canonicalize-vector-extract.ll (+10)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index e29fb869686ca0b..00505b5edb58dca 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -2997,24 +2997,26 @@ 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();
       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 (dyn_cast<ScalableVectorType>(VecTy) || 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
 ; ============================================================================ ;

[david-arm]

LGTM! I suggested a possible tweak you could make before landing the patch, but only if you think it's worth it. :)