[PATCH] D102404: [InstCombine] Add instcombine fold for extractelement + splat for scalable vectors

[Caroline via Phabricator via llvm-commits]

CarolineConcatto created this revision.
Herald added a subscriber: hiraditya.
CarolineConcatto requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

This patch allows that scalable vector can also use the fold that already
exists for fixed vector, only when the lane index is lower than the minimum
number of elements of the vector.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D102404

Files:
  llvm/lib/Analysis/VectorUtils.cpp
  llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll
  llvm/test/Transforms/InstCombine/vscale_extractelement.ll


Index: llvm/test/Transforms/InstCombine/vscale_extractelement.ll
===================================================================
--- llvm/test/Transforms/InstCombine/vscale_extractelement.ll
+++ llvm/test/Transforms/InstCombine/vscale_extractelement.ll
@@ -55,13 +55,9 @@
   ret i8 %r
 }
 
-; TODO: Instcombine could optimize to return %v.
 define i32 @extractelement_shuffle_in_range(i32 %v) {
 ; CHECK-LABEL: @extractelement_shuffle_in_range(
-; CHECK-NEXT:    [[IN:%.*]] = insertelement <vscale x 4 x i32> undef, i32 [[V:%.*]], i32 0
-; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i32 1
-; CHECK-NEXT:    ret i32 [[R]]
+; CHECK-NEXT:    ret i32 %v
 ;
   %in = insertelement <vscale x 4 x i32> undef, i32 %v, i32 0
   %splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
Index: llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll
===================================================================
--- llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll
+++ llvm/test/Transforms/InstCombine/vscale_extractelement-inseltpoison.ll
@@ -55,13 +55,9 @@
   ret i8 %r
 }
 
-; TODO: Instcombine could optimize to return %v.
 define i32 @extractelement_shuffle_in_range(i32 %v) {
 ; CHECK-LABEL: @extractelement_shuffle_in_range(
-; CHECK-NEXT:    [[IN:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[V:%.*]], i32 0
-; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i32 1
-; CHECK-NEXT:    ret i32 [[R]]
+; CHECK-NEXT:    ret i32 %v
 ;
   %in = insertelement <vscale x 4 x i32> poison, i32 %v, i32 0
   %splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
Index: llvm/lib/Analysis/VectorUtils.cpp
===================================================================
--- llvm/lib/Analysis/VectorUtils.cpp
+++ llvm/lib/Analysis/VectorUtils.cpp
@@ -301,16 +301,18 @@
   }
 
   ShuffleVectorInst *SVI = dyn_cast<ShuffleVectorInst>(V);
-  // Restrict the following transformation to fixed-length vector.
-  if (SVI && isa<FixedVectorType>(SVI->getType())) {
-    unsigned LHSWidth =
-        cast<FixedVectorType>(SVI->getOperand(0)->getType())->getNumElements();
-    int InEl = SVI->getMaskValue(EltNo);
-    if (InEl < 0)
-      return UndefValue::get(VTy->getElementType());
-    if (InEl < (int)LHSWidth)
-      return findScalarElement(SVI->getOperand(0), InEl);
-    return findScalarElement(SVI->getOperand(1), InEl - LHSWidth);
+  if (SVI && isa<VectorType>(SVI->getType())) {
+    ElementCount LHS =
+        cast<VectorType>(SVI->getOperand(0)->getType())->getElementCount();
+    unsigned LHSWidth = LHS.getKnownMinValue();
+    if ((EltNo < LHSWidth && LHS.isScalable()) || !LHS.isScalable()) {
+      int InEl = SVI->getMaskValue(EltNo);
+      if (InEl < 0)
+        return UndefValue::get(VTy->getElementType());
+      if (InEl < (int)LHSWidth)
+        return findScalarElement(SVI->getOperand(0), InEl);
+      return findScalarElement(SVI->getOperand(1), InEl - LHSWidth);
+    }
   }
 
   // Extract a value from a vector add operation with a constant zero.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D102404.345146.patch
Type: text/x-patch
Size: 3452 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20210513/dc0a84c7/attachment.bin>