[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.

Author: kmclaughlin-arm

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);

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
 ; ============================================================================ ;