[VectorCombine] Enable transform 'foldSingleElementStore' for scalable vector types

[benshi001]

Summary

The transform 'foldSingleElementStore' can be applied to scalable
vector types if the index is less than the minimum number of elements.

[dmgreen]

Does this need to be limited to LE? My understanding is that BE and LE vectors would store the elements in the same places, even for scalable vectors. (i.e. it's like an array, where the 0th element ends up first in memory).

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

49

Can you remove the unrelated formatting changes.

[benshi001]

In D157676#4587470, @dmgreen wrote:

Does this need to be limited to LE? My understanding is that BE and LE vectors would store the elements in the same places, even for scalable vectors. (i.e. it's like an array, where the 0th element ends up first in memory).

Thanks. I have checked with AArch64's SVE standard and RISCV V-extension specification, LE/BE are just concepts of byte order within a scalar element, not the order of elements inside a vector.

So I applied my optimization to scalable vector types without check of BE/LE.

Thanks.

[dmgreen]

Thanks. That matches my understanding of llvm vectors too. LGTM

[nikic]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

1023–1026

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

 STATISTIC(NumVecCmpBO, "Number of vector compare + binop formed");
 STATISTIC(NumShufOfBitcast, "Number of shuffles moved after bitcast");
 STATISTIC(NumScalarBO, "Number of scalar binops formed");
 STATISTIC(NumScalarCmp, "Number of scalar compares formed");
 
 static cl::opt<bool> DisableVectorCombine(
     "disable-vector-combine", cl::init(false), cl::Hidden,
     cl::desc("Disable all vector combine transforms"));
 
 static cl::opt<bool> DisableBinopExtractShuffle(
     "disable-binop-extract-shuffle", cl::init(false), cl::Hidden,
     cl::desc("Disable binop extract to shuffle transforms"));
 
 static cl::opt<unsigned> MaxInstrsToScan(
     "vector-combine-max-scan-instrs", cl::init(30), cl::Hidden,
     cl::desc("Max number of instructions to scan for vector combining."));
 
void freeze(IRBuilder<> &Builder, Instruction &UserI) {
 
     ToFreeze = nullptr;
   }
 };
 } // namespace
 
 /// Check if it is legal to scalarize a memory access to \p VecTy at index \p
 /// Idx. \p Idx must access a valid vector element.
-static ScalarizationResult canScalarizeAccess(FixedVectorType *VecTy,
-                                              Value *Idx, Instruction *CtxI,
+static ScalarizationResult canScalarizeAccess(VectorType *VecTy, Value *Idx,
+                                              Instruction *CtxI,
                                               AssumptionCache &AC,
                                               const DominatorTree &DT) {
+  // We do checks for both fixed vector types and scalable vector types.
+  // This is the number of elements of fixed vector types,
+  // or the minium number of elements of scalable vector types.
+  uint64_t NumElements = VecTy->getElementCount().getKnownMinValue();
+
   if (auto *C = dyn_cast<ConstantInt>(Idx)) {
-    if (C->getValue().ult(VecTy->getNumElements()))
+    if (C->getValue().ult(NumElements))
       return ScalarizationResult::safe();
     return ScalarizationResult::unsafe();
   }
 
   unsigned IntWidth = Idx->getType()->getScalarSizeInBits();
   APInt Zero(IntWidth, 0);
-  APInt MaxElts(IntWidth, VecTy->getNumElements());
+  APInt MaxElts(IntWidth, NumElements);
   ConstantRange ValidIndices(Zero, MaxElts);
   ConstantRange IdxRange(IntWidth, true);
 
   if (isGuaranteedNotToBePoison(Idx, &AC)) {
     if (ValidIndices.contains(computeConstantRange(Idx, /* ForSigned */ false,
                                                    true, &AC, CtxI, &DT)))
       return ScalarizationResult::safe();
     return ScalarizationResult::unsafe();

 //   %1 = insertelement <4 x i32> %0, i32 %b, i32 1
 //   store <4 x i32> %1, <4 x i32>* %a
 // to:
 //   %0 = bitcast <4 x i32>* %a to i32*
 //   %1 = getelementptr inbounds i32, i32* %0, i64 0, i64 1
 //   store i32 %b, i32* %1
 bool VectorCombine::foldSingleElementStore(Instruction &I) {
   auto *SI = cast<StoreInst>(&I);
-  if (!SI->isSimple() ||
-      !isa<FixedVectorType>(SI->getValueOperand()->getType()))
+  if (!SI->isSimple() || !isa<VectorType>(SI->getValueOperand()->getType()))
     return false;
 
   // TODO: Combine more complicated patterns (multiple insert) by referencing
   // TargetTransformInfo.
   Instruction *Source;
   Value *NewElement;
   Value *Idx;
   if (!match(SI->getValueOperand(),
              m_InsertElt(m_Instruction(Source), m_Value(NewElement),
                          m_Value(Idx))))
     return false;
 
   if (auto *Load = dyn_cast<LoadInst>(Source)) {
-    auto VecTy = cast<FixedVectorType>(SI->getValueOperand()->getType());
+    auto VecTy = cast<VectorType>(SI->getValueOperand()->getType());
     const DataLayout &DL = I.getModule()->getDataLayout();
     Value *SrcAddr = Load->getPointerOperand()->stripPointerCasts();
     // Don't optimize for atomic/volatile load or store. Ensure memory is not
     // modified between, vector type matches store size, and index is inbounds.
     if (!Load->isSimple() || Load->getParent() != SI->getParent() ||
         !DL.typeSizeEqualsStoreSize(Load->getType()) ||
         SrcAddr != SI->getPointerOperand()->stripPointerCasts())
       return false;

llvm/test/Transforms/VectorCombine/load-insert-store.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -S -passes=vector-combine -data-layout=e < %s | FileCheck %s --check-prefixes=CHECK,LE
-; RUN: opt -S -passes=vector-combine -data-layout=E < %s | FileCheck %s --check-prefixes=CHECK,BE
+; RUN: opt -S -passes=vector-combine -data-layout=e < %s | FileCheck %s
+; RUN: opt -S -passes=vector-combine -data-layout=E < %s | FileCheck %s
 
 define void @insert_store(ptr %q, i8 zeroext %s) {
 ; CHECK-LABEL: @insert_store(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <16 x i8>, ptr [[Q:%.*]], i32 0, i32 3
 ; CHECK-NEXT:    store i8 [[S:%.*]], ptr [[TMP0]], align 1
 ; CHECK-NEXT:    ret void
 ;
entry:
   %vecins = insertelement <8 x i16> %0, i16 %s, i32 9
   store <8 x i16> %vecins, ptr %q
   ret void
 }
 
 define void @insert_store_vscale(ptr %q, i16 zeroext %s) {
 ; CHECK-LABEL: @insert_store_vscale(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = load <vscale x 8 x i16>, ptr [[Q:%.*]], align 16
-; CHECK-NEXT:    [[VECINS:%.*]] = insertelement <vscale x 8 x i16> [[TMP0]], i16 [[S:%.*]], i32 3
-; CHECK-NEXT:    store <vscale x 8 x i16> [[VECINS]], ptr [[Q]], align 16
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <vscale x 8 x i16>, ptr [[Q:%.*]], i32 0, i32 3
+; CHECK-NEXT:    store i16 [[S:%.*]], ptr [[TMP0]], align 2
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %0 = load <vscale x 8 x i16>, ptr %q
   %vecins = insertelement <vscale x 8 x i16> %0, i16 %s, i32 3
   store <vscale x 8 x i16> %vecins, ptr %q
   ret void
 }

 
 ; To verify the index is not a constant but valid by assume,
 ; for scalable vector types.
 define void @insert_store_vscale_nonconst_index_known_valid_by_assume(ptr %q, i8 zeroext %s, i32 %idx) {
 ; CHECK-LABEL: @insert_store_vscale_nonconst_index_known_valid_by_assume(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[IDX:%.*]], 4
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
-; CHECK-NEXT:    [[TMP0:%.*]] = load <vscale x 16 x i8>, ptr [[Q:%.*]], align 16
-; CHECK-NEXT:    [[VECINS:%.*]] = insertelement <vscale x 16 x i8> [[TMP0]], i8 [[S:%.*]], i32 [[IDX]]
-; CHECK-NEXT:    store <vscale x 16 x i8> [[VECINS]], ptr [[Q]], align 16
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <vscale x 16 x i8>, ptr [[Q:%.*]], i32 0, i32 [[IDX]]
+; CHECK-NEXT:    store i8 [[S:%.*]], ptr [[TMP0]], align 1
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %cmp = icmp ult i32 %idx, 4
   call void @llvm.assume(i1 %cmp)
   %0 = load <vscale x 16 x i8>, ptr %q
   %vecins = insertelement <vscale x 16 x i8> %0, i8 %s, i32 %idx
   store <vscale x 16 x i8> %vecins, ptr %q
entry:
   ret void
 }
 
 ; To verify the index is not a constant but valid by and,
 ; for scalable vector types.
 define void @insert_store_vscale_nonconst_index_known_noundef_and_valid_by_and(ptr %q, i8 zeroext %s, i32 noundef %idx) {
 ; CHECK-LABEL: @insert_store_vscale_nonconst_index_known_noundef_and_valid_by_and(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = load <vscale x 16 x i8>, ptr [[Q:%.*]], align 16
 ; CHECK-NEXT:    [[IDX_CLAMPED:%.*]] = and i32 [[IDX:%.*]], 7
-; CHECK-NEXT:    [[VECINS:%.*]] = insertelement <vscale x 16 x i8> [[TMP0]], i8 [[S:%.*]], i32 [[IDX_CLAMPED]]
-; CHECK-NEXT:    store <vscale x 16 x i8> [[VECINS]], ptr [[Q]], align 16
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <vscale x 16 x i8>, ptr [[Q:%.*]], i32 0, i32 [[IDX_CLAMPED]]
+; CHECK-NEXT:    store i8 [[S:%.*]], ptr [[TMP0]], align 1
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %0 = load <vscale x 16 x i8>, ptr %q
   %idx.clamped = and i32 %idx, 7
   %vecins = insertelement <vscale x 16 x i8> %0, i8 %s, i32 %idx.clamped
   store <vscale x 16 x i8> %vecins, ptr %q
   ret void
entry:
   ret void
 }
 
 ; To verify the index is not a constant but valid by urem,
 ; for scalable vector types.
 define void @insert_store_vscale_nonconst_index_known_noundef_and_valid_by_urem(ptr %q, i8 zeroext %s, i32 noundef %idx) {
 ; CHECK-LABEL: @insert_store_vscale_nonconst_index_known_noundef_and_valid_by_urem(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = load <vscale x 16 x i8>, ptr [[Q:%.*]], align 16
 ; CHECK-NEXT:    [[IDX_CLAMPED:%.*]] = urem i32 [[IDX:%.*]], 16
-; CHECK-NEXT:    [[VECINS:%.*]] = insertelement <vscale x 16 x i8> [[TMP0]], i8 [[S:%.*]], i32 [[IDX_CLAMPED]]
-; CHECK-NEXT:    store <vscale x 16 x i8> [[VECINS]], ptr [[Q]], align 16
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <vscale x 16 x i8>, ptr [[Q:%.*]], i32 0, i32 [[IDX_CLAMPED]]
+; CHECK-NEXT:    store i8 [[S:%.*]], ptr [[TMP0]], align 1
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %0 = load <vscale x 16 x i8>, ptr %q
   %idx.clamped = urem i32 %idx, 16
   %vecins = insertelement <vscale x 16 x i8> %0, i8 %s, i32 %idx.clamped
   store <vscale x 16 x i8> %vecins, ptr %q
   ret void
bb:
   %i35 = lshr i32 %i34, 1
   %i36 = insertelement <16 x i8> %i4, i8 %arg3, i32 3
   store <16 x i8> %i36, ptr %arg2, align 16
   ret i32 %i35
 }
 
 declare i32 @bar(i32, i1) readonly
 declare double @llvm.log2.f64(double)
-;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
-; BE: {{.*}}
-; LE: {{.*}}