================
@@ -5191,47 +5192,95 @@ bool SROA::presplitLoadsAndStores(AllocaInst &AI, 
AllocaSlices &AS) {
 /// promoted.
 AllocaInst *SROA::rewritePartition(AllocaInst &AI, AllocaSlices &AS,
                                    Partition &P) {
+  const DataLayout &DL = AI.getDataLayout();
   // Try to compute a friendly type for this partition of the alloca. This
   // 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;
-  const DataLayout &DL = AI.getDataLayout();
-  unsigned VScale = AI.getFunction()->getVScaleValue();
-
-  std::pair<Type *, IntegerType *> CommonUseTy =
-      findCommonType(P.begin(), P.end(), P.endOffset());
-  // Do all uses operate on the same type?
-  if (CommonUseTy.first) {
-    TypeSize CommonUseSize = DL.getTypeAllocSize(CommonUseTy.first);
-    if (CommonUseSize.isFixed() && CommonUseSize.getFixedValue() >= P.size())
-      SliceTy = CommonUseTy.first;
-  }
-  // If not, can we find an appropriate subtype in the original allocated type?
-  if (!SliceTy)
+  //
+  // Returns a tuple with the following elements:
+  //   - PartitionType: The computed type for this partition.
+  //   - IsIntegerWideningViable: True if integer widening promotion is used.
+  //   - VectorType: The vector type if vector promotion is used, otherwise
+  //   nullptr.
+  auto SelectPartitionTy = [&]() -> std::tuple<Type *, bool, VectorType *> {
+    // First check if the partition is viable for vetor promotion.
+    //
+    // We prefer vector promotion over integer widening promotion when:
+    // - The vector element type is a floating-point type.
+    // - All the loads/stores to the alloca are vector loads/stores to the
+    // entire alloca.
+    //
+    // Otherwise when there is a integer vector with mixed
+    // loads/stores we prefer integer widening promotion because it's more
+    // likely the user is doing bitwise arithmetic and we generate better code.
+    VectorType *VecTy =
+        isVectorPromotionViable(P, DL, AI.getFunction()->getVScaleValue());
+    // If the vector element type is a floating-point type, we prefer vector
+    // promotion. If the vector has one element, let the below code select
+    // whether we promote with the vector or scalar.
+    if (VecTy && VecTy->getElementType()->isFloatingPointTy() &&
+        VecTy->getElementCount().getFixedValue() > 1)
+      return {VecTy, false, VecTy};
+
+    // Check if there is a common type that all slices of the partition use 
that
+    // spans the partition.
+    auto [CommonUseTy, LargestIntTy] =
+        findCommonType(P.begin(), P.end(), P.endOffset());
+    if (CommonUseTy) {
+      TypeSize CommonUseSize = DL.getTypeAllocSize(CommonUseTy);
+      if (CommonUseSize.isFixed() &&
+          CommonUseSize.getFixedValue() >= P.size()) {
+        // We prefer vector promotion here because if vector promotion is 
viable
+        // and there is a common type used, then it implies the second listed
+        // condition for prefering vector promotion is true.
+        if (VecTy)
+          return {VecTy, false, VecTy};
+        return {CommonUseTy, isIntegerWideningViable(P, CommonUseTy, DL),
+                nullptr};
+      }
+    }
+
+    // Can we find an appropriate subtype in the original allocated
+    // type?
     if (Type *TypePartitionTy = getTypePartition(DL, AI.getAllocatedType(),
-                                                 P.beginOffset(), P.size()))
-      SliceTy = TypePartitionTy;
+                                                 P.beginOffset(), P.size())) {
+      // If the partition is an integer array that can be spanned by a legal
+      // integer type, prefer to represent it as a legal integer type because
+      // it's more likely to be promotable.
+      if (TypePartitionTy->isArrayTy() &&
+          TypePartitionTy->getArrayElementType()->isIntegerTy() &&
+          DL.isLegalInteger(P.size() * 8))
+        TypePartitionTy = Type::getIntNTy(*C, P.size() * 8);
+      // There was no common type used, so we prefer integer widening 
promotion.
+      if (isIntegerWideningViable(P, TypePartitionTy, DL))
+        return {TypePartitionTy, true, nullptr};
+      if (VecTy)
+        return {VecTy, false, VecTy};
+      // If we couldn't promotion with TypePartitionTy, try with the largest
+      // integer type used.
+      if (LargestIntTy &&
+          DL.getTypeAllocSize(LargestIntTy).getFixedValue() >= P.size() &&
+          isIntegerWideningViable(P, LargestIntTy, DL))
+        return {LargestIntTy, true, nullptr};
+
+      // Fallback to TypePartitionTy and we probably won't promote.
+      return {TypePartitionTy, false, nullptr};
+    }
 
-  // If still not, can we use the largest bitwidth integer type used?
-  if (!SliceTy && CommonUseTy.second)
-    if (DL.getTypeAllocSize(CommonUseTy.second).getFixedValue() >= P.size())
-      SliceTy = CommonUseTy.second;
-  if ((!SliceTy || (SliceTy->isArrayTy() &&
-                    SliceTy->getArrayElementType()->isIntegerTy())) &&
-      DL.isLegalInteger(P.size() * 8)) {
-    SliceTy = Type::getIntNTy(*C, P.size() * 8);
-  }
+    // Select the largest integer type used if it spans the partition.
+    if (LargestIntTy &&
+        DL.getTypeAllocSize(LargestIntTy).getFixedValue() >= P.size())
+      return {LargestIntTy, false, nullptr};
 
-  if (!SliceTy)
-    SliceTy = ArrayType::get(Type::getInt8Ty(*C), P.size());
-  assert(DL.getTypeAllocSize(SliceTy).getFixedValue() >= P.size());
+    // Select a legal integer type if it spans the partition.
+    if (DL.isLegalInteger(P.size() * 8))
+      return {Type::getIntNTy(*C, P.size() * 8), false, nullptr};
----------------
YonahGoldberg wrote:

You want to move both of these integer selection cases above the 
`getTypePartition()` call?

I don't think that would work. I think there are many cases where you'd prefer 
to select the type partition over an integer, but maybe there's something I'm 
missing.

But in any case in the original code the getTypePartition() type is selected 
with higher priority so I'd say that's really beyond the scope of this PR where 
I was just trying to reorganize so I could capture the edge case I encountered.

https://github.com/llvm/llvm-project/pull/167771
_______________________________________________
cfe-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to