[SROA] Clean up rewritePartition type selection process

[YonahGoldberg]

This change reverts 2572512, which landed on Aug 8, 2022. This change addressed the problem that if you have IR that looks something like:

%alloca = alloca <4 x float>
store <4 x float> %data, ptr %alloca
%load = load half, ptr %alloca

getCommonType would return <4 x float> because the load half isn't to the entire partition, so we skip the first getTypePartition check. 2572512 added a later check that sees that <4 x float> is not vector promotable because of the load half, and then calls getTypePartition, which changes the sliceTy to < 8 x half>, which is vector promotable because the store can be changed to store <8 x half>. So we set the sliceTy to <8 x half>, we can promote the alloca, and everyone is happy.

This code became unnecessary after 529eafd, which landed ~3 months later, which fixes the issue in a different way. isVectorPromotionViable was already smart enough to try <8 x half> as a type candidate because it sees the load half. However, this candidate didn't work because it conflicts with store <4 x float>. This commit added logic to try integer-ifying candidates if there is no common type. So the <8 x half> candidate gets converted to <8 x i16>, which works because we can convert the alloca to alloca <8 x i16> and the load to load i16, allowing promotion.

After 529eafd, the original commit is pointless. It tries to refine the SliceTy, but if isVectorPromotionViable succeeds, it returns a new type to use and we will ignore the SliceTy.

This change is my first patch to try to simplify the type selection process in rewritePartition. I had some other ideas that I tried in #167771 and #168796, but they need refinement.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yonah Goldberg (YonahGoldberg)

Changes

This change reverts 2572512, which landed on Aug 8, 2022. This change addressed the problem that if you have IR that looks something like:

%alloca = alloca <4 x float>
store <4 x float> %data, ptr %alloca
%load = load half, ptr %alloca

getCommonType would return <4 x float> because the load half isn't to the entire partition, so we skip the first getTypePartition check. 2572512 added a later check that sees that <4 x float> is not vector promotable because of the load half, and then calls getTypePartition, which changes the sliceTy to < 8 x half>, which is vector promotable because the store can be changed to store <8 x half>. So we set the sliceTy to <8 x half>, we can promote the alloca, and everyone is happy.

This code became unnecessary after 529eafd, which landed ~3 months later, which fixes the issue in a different way. isVectorPromotionViable was already smart enough to try <8 x half> as a type candidate because it sees the load half. However, this candidate didn't work because it conflicts with store <4 x float>. This commit added logic to try integer-ifying candidates if there is no common type. So the <8 x half> candidate gets converted to <8 x i16>, which works because we can convert the alloca to alloca <8 x i16> and the load to load i16, allowing promotion.

After 529eafd, the original commit is pointless. It tries to refine the SliceTy, but if isVectorPromotionViable succeeds, it returns a new type to use and we will ignore the SliceTy.

This change is my first patch to try to simplify the type selection process in rewritePartition. A few other issues I've noticed are. I had some other ideas that I tried in #167771 and #168796, but they need refinement.

---

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

1 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/SROA.cpp (+25-51)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 5c60fad6f91aa..31ecb2391412a 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -2178,35 +2178,6 @@ static bool isVectorPromotionViableForSlice(Partition &P, const Slice &S,
   return true;
 }
 
-/// Test whether a vector type is viable for promotion.
-///
-/// This implements the necessary checking for \c checkVectorTypesForPromotion
-/// (and thus isVectorPromotionViable) over all slices of the alloca for the
-/// given VectorType.
-static bool checkVectorTypeForPromotion(Partition &P, VectorType *VTy,
-                                        const DataLayout &DL, unsigned VScale) {
-  uint64_t ElementSize =
-      DL.getTypeSizeInBits(VTy->getElementType()).getFixedValue();
-
-  // While the definition of LLVM vectors is bitpacked, we don't support sizes
-  // that aren't byte sized.
-  if (ElementSize % 8)
-    return false;
-  assert((DL.getTypeSizeInBits(VTy).getFixedValue() % 8) == 0 &&
-         "vector size not a multiple of element size?");
-  ElementSize /= 8;
-
-  for (const Slice &S : P)
-    if (!isVectorPromotionViableForSlice(P, S, VTy, ElementSize, DL, VScale))
-      return false;
-
-  for (const Slice *S : P.splitSliceTails())
-    if (!isVectorPromotionViableForSlice(P, *S, VTy, ElementSize, DL, VScale))
-      return false;
-
-  return true;
-}
-
 /// Test whether any vector type in \p CandidateTys is viable for promotion.
 ///
 /// This implements the necessary checking for \c isVectorPromotionViable over
@@ -2291,11 +2262,30 @@ checkVectorTypesForPromotion(Partition &P, const DataLayout &DL,
            std::numeric_limits<unsigned short>::max();
   });
 
-  for (VectorType *VTy : CandidateTys)
-    if (checkVectorTypeForPromotion(P, VTy, DL, VScale))
-      return VTy;
+  // Find a vector type viable for promotion by iterating over all slices.
+  auto *VTy = llvm::find_if(CandidateTys, [&](VectorType *VTy) -> bool {
+    uint64_t ElementSize =
+        DL.getTypeSizeInBits(VTy->getElementType()).getFixedValue();
 
-  return nullptr;
+    // While the definition of LLVM vectors is bitpacked, we don't support sizes
+    // that aren't byte sized.
+    if (ElementSize % 8)
+      return false;
+    assert((DL.getTypeSizeInBits(VTy).getFixedValue() % 8) == 0 &&
+           "vector size not a multiple of element size?");
+    ElementSize /= 8;
+
+    for (const Slice &S : P)
+      if (!isVectorPromotionViableForSlice(P, S, VTy, ElementSize, DL, VScale))
+        return false;
+
+    for (const Slice *S : P.splitSliceTails())
+      if (!isVectorPromotionViableForSlice(P, *S, VTy, ElementSize, DL, VScale))
+        return false;
+
+    return true;
+  });
+  return VTy != CandidateTys.end() ? *VTy : nullptr;
 }
 
 static VectorType *createAndCheckVectorTypesForPromotion(
@@ -5213,7 +5203,6 @@ AllocaInst *SROA::rewritePartition(AllocaInst &AI, AllocaSlices &AS,
   // won't always succeed, in which case we fall back to a legal integer type
   // or an i8 array of an appropriate size.
   Type *SliceTy = nullptr;
-  VectorType *SliceVecTy = nullptr;
   const DataLayout &DL = AI.getDataLayout();
   unsigned VScale = AI.getFunction()->getVScaleValue();
 
@@ -5222,10 +5211,8 @@ AllocaInst *SROA::rewritePartition(AllocaInst &AI, AllocaSlices &AS,
   // Do all uses operate on the same type?
   if (CommonUseTy.first) {
     TypeSize CommonUseSize = DL.getTypeAllocSize(CommonUseTy.first);
-    if (CommonUseSize.isFixed() && CommonUseSize.getFixedValue() >= P.size()) {
+    if (CommonUseSize.isFixed() && CommonUseSize.getFixedValue() >= P.size())
       SliceTy = CommonUseTy.first;
-      SliceVecTy = dyn_cast<VectorType>(SliceTy);
-    }
   }
   // If not, can we find an appropriate subtype in the original allocated type?
   if (!SliceTy)
@@ -5235,27 +5222,14 @@ AllocaInst *SROA::rewritePartition(AllocaInst &AI, AllocaSlices &AS,
 
   // If still not, can we use the largest bitwidth integer type used?
   if (!SliceTy && CommonUseTy.second)
-    if (DL.getTypeAllocSize(CommonUseTy.second).getFixedValue() >= P.size()) {
+    if (DL.getTypeAllocSize(CommonUseTy.second).getFixedValue() >= P.size())
       SliceTy = CommonUseTy.second;
-      SliceVecTy = dyn_cast<VectorType>(SliceTy);
-    }
   if ((!SliceTy || (SliceTy->isArrayTy() &&
                     SliceTy->getArrayElementType()->isIntegerTy())) &&
       DL.isLegalInteger(P.size() * 8)) {
     SliceTy = Type::getIntNTy(*C, P.size() * 8);
   }
 
-  // If the common use types are not viable for promotion then attempt to find
-  // another type that is viable.
-  if (SliceVecTy && !checkVectorTypeForPromotion(P, SliceVecTy, DL, VScale))
-    if (Type *TypePartitionTy = getTypePartition(DL, AI.getAllocatedType(),
-                                                 P.beginOffset(), P.size())) {
-      VectorType *TypePartitionVecTy = dyn_cast<VectorType>(TypePartitionTy);
-      if (TypePartitionVecTy &&
-          checkVectorTypeForPromotion(P, TypePartitionVecTy, DL, VScale))
-        SliceTy = TypePartitionTy;
-    }
-
   if (!SliceTy)
     SliceTy = ArrayType::get(Type::getInt8Ty(*C), P.size());
   assert(DL.getTypeAllocSize(SliceTy).getFixedValue() >= P.size());

[github-actions]

🐧 Linux x64 Test Results

• 186457 tests passed
• 4867 tests skipped

[YonahGoldberg]

@dtcxzyw @nikic @vangthao95 @LebedevRI could I get a review?

[vangthao95]

LGTM.