[SLP] Fix for a bug in jumbled memory access vectorization introduced in r293386

[• ashahid]

Summary

The bug is due to absence of in order uses of scalars which needs to be available for VectorizeTree() API. This API uses it for proper mask computation to be used in "shufflevector" IR.

The fix is to compute the mask for out of order memory accesses while building the vectorizable tree instead of actual vectorization of vectorizable tree.

[mkuper]

This is getting a bit hairy, I think - we have too many valuelists flying around.

How about something like this:
When we sort the vectors, we can also use the sort to generate the corresponding shuffle mask (instead of recomputing it in O(|VL|^2) time in line 2623 - 2634), by sorting pairs of <Value, Index>.
We can then keep an optional shuffle mask instead of the NeedToShuffle bit as part of the tree node, and use it to get the correctly ordered list of scalars back from VectorizableTree[0].Scalars when we need to.

What do you think? My gut feeling is that keeping the mask around is worth it just to avoid the N^2 loop in 2623, and the fix to the bug falls out of it.

[• ashahid]

In D30159#681835, @mkuper wrote:

This is getting a bit hairy, I think - we have too many valuelists flying around.

How about something like this:
When we sort the vectors, we can also use the sort to generate the corresponding shuffle mask (instead of recomputing it in O(|VL|^2) time in line 2623 - 2634), by sorting pairs of <Value, Index>.

The basic idea to keep the mask computation at the current place was because we wanted to compute the mask, (N^2 loop), only if vectorization is found to be beneficial.

OTOH if we do mask computation while sorting memory accesses, the utility of SortMemAccesses API may be reduced, however this is not a bigger problem and can be
worked around by keeping a bit to decide whether to compute the mask or not.

We can then keep an optional shuffle mask instead of the NeedToShuffle bit as part of the tree node, and use it to get the correctly ordered list of scalars back from VectorizableTree[0].Scalars when we need to.

What do you think? My gut feeling is that keeping the mask around is worth it just to avoid the N^2 loop in 2623, and the fix to the bug falls out of it.

Although I am not an expert, I think keeping another ValueList for unordered scalars is a not a bad design. However I am open to other views.

Regards,
Shahid

[mkuper]

In D30159#681901, @ashahid wrote:

In D30159#681835, @mkuper wrote:

This is getting a bit hairy, I think - we have too many valuelists flying around.

How about something like this:
When we sort the vectors, we can also use the sort to generate the corresponding shuffle mask (instead of recomputing it in O(|VL|^2) time in line 2623 - 2634), by sorting pairs of <Value, Index>.

The basic idea to keep the mask computation at the current place was because we wanted to compute the mask, (N^2 loop), only if vectorization is found to be beneficial.

Yes, but we can avoid the O(N^2) computation altogether, right? I mean, doing that doesn't add any cost to the sort. Or am I missing something?

OTOH if we do mask computation while sorting memory accesses, the utility of SortMemAccesses API may be reduced, however this is not a bigger problem and can be
worked around by keeping a bit to decide whether to compute the mask or not.

Don't we sort before we know whether we're going to need the mask?

We can then keep an optional shuffle mask instead of the NeedToShuffle bit as part of the tree node, and use it to get the correctly ordered list of scalars back from VectorizableTree[0].Scalars when we need to.

What do you think? My gut feeling is that keeping the mask around is worth it just to avoid the N^2 loop in 2623, and the fix to the bug falls out of it.

Although I am not an expert, I think keeping another ValueList for unordered scalars is a not a bad design. However I am open to other views.

It seems odd to have both the sorted and the unsorted list. The other thing that bothers me is that we have some invariants which don't seem to be well-defined. Is the guarantee that "Scalars" is ordered to have the "best" order to allow vectorization, while InOrdScalars is ordered according to the original IR order? Also, preserving "program order" is kind of meaningless. The fact we care about order at all is just an artifact of the way we build the tree. What we really care about is mapping each scalar leaf to the vector lane it eventually ends up in, which is precisely the shuffle mask.

Anyway, I'm also open to other opinions, if more people care about this.
What I'm really trying to do here is to make the code make a bit more sense - right now, it makes almost zero sense to me. (This isn't the fault of your patch, I'm talking about the general state of the SLP vectorizer).

[• ashahid]

In D30159#681919, @mkuper wrote:

In D30159#681901, @ashahid wrote:

In D30159#681835, @mkuper wrote:

This is getting a bit hairy, I think - we have too many valuelists flying around.

How about something like this:
When we sort the vectors, we can also use the sort to generate the corresponding shuffle mask (instead of recomputing it in O(|VL|^2) time in line 2623 - 2634), by sorting pairs of <Value, Index>.

The basic idea to keep the mask computation at the current place was because we wanted to compute the mask, (N^2 loop), only if vectorization is found to be beneficial.

Yes, but we can avoid the O(N^2) computation altogether, right? I mean, doing that doesn't add any cost to the sort. Or am I missing something?

OTOH if we do mask computation while sorting memory accesses, the utility of SortMemAccesses API may be reduced, however this is not a bigger problem and can be
worked around by keeping a bit to decide whether to compute the mask or not.

Don't we sort before we know whether we're going to need the mask?

We can then keep an optional shuffle mask instead of the NeedToShuffle bit as part of the tree node, and use it to get the correctly ordered list of scalars back from VectorizableTree[0].Scalars when we need to.

What do you think? My gut feeling is that keeping the mask around is worth it just to avoid the N^2 loop in 2623, and the fix to the bug falls out of it.

Although I am not an expert, I think keeping another ValueList for unordered scalars is a not a bad design. However I am open to other views.

It seems odd to have both the sorted and the unsorted list. The other thing that bothers me is that we have some invariants which don't seem to be well-defined. Is the guarantee that "Scalars" is ordered to have the "best" order to allow vectorization, while InOrdScalars is ordered according to the original IR order? Also, preserving "program order" is kind of meaningless. The fact we care about order at all is just an artifact of the way we build the tree. What we really care about is mapping each scalar leaf to the vector lane it eventually ends up in, which is precisely the shuffle mask.

Anyway, I'm also open to other opinions, if more people care about this.
What I'm really trying to do here is to make the code make a bit more sense - right now, it makes almost zero sense to me. (This isn't the fault of your patch, I'm talking about the general state of the SLP vectorizer).

Thanks Michael, my bad, I completely missed that we can avoid O(N^2) mask computation altogether. Moreover, we don't need to have extra VL and NeedToShuffle.
I will update the patch accordingly.

[• ashahid]

Updated the patch to incorporate the review comments.

[mkuper]

lib/Analysis/LoopAccessAnalysis.cpp

1062 ↗

(On Diff #89610)

The API seems a bit odd - having to pass a SmallVector when you know you don't want it to be filled doesn't look right.
I'd prefer something like passing in a pointer, and filling the vector in if it's not null (or passing in an Optional<> reference to a similar effect.)

1098 ↗

(On Diff #89610)

I was thinking along the lines of directly sorting an array of tuples <Offset, Value*, Index> - and then you can pull the mask out of it.
But the that may turn out uglier than what you have here. Your call.

1111 ↗

(On Diff #89610)

Redundant braces.

lib/Transforms/Vectorize/SLPVectorizer.cpp

2640

This seems unrelated.
Should it have been here before as well?

[• ashahid]

lib/Analysis/LoopAccessAnalysis.cpp

1062 ↗

(On Diff #89610)

Ok

1098 ↗

(On Diff #89610)

With tuple, I think we still need 2 std::sort() call. I am opting for the simpler than uglier.

lib/Transforms/Vectorize/SLPVectorizer.cpp

2640

Earlier this path was not tested properly but thanks to this bug we found.

As in this particular case the vectorized value is being used from the tree entry data instead of the return value, if we don't update E->VectorizedValue "shufflvector" IR will not be used, rather vectorized load IR will be used which is not the intention.

[• ashahid]

Updated the path to incorporate review comments.

[mkuper]

LGTM, with some style comments inline.

lib/Analysis/LoopAccessAnalysis.cpp

1097 ↗

(On Diff #89974)

IIRC, you can do something like "SmallVector<unsigned, 4> UseOrder(VL.size())", and then "UseOrder[i] = i".
This is slightly shorter, and maybe a bit prettier, but feel free to leave as is if you prefer it.

1115 ↗

(On Diff #89974)

I think I would group everything that has to do with Mask together - something like:

if (Mask) {
  Mask->reserve(VL.size());
  for (unsigned i = 0; i < VL.size(); i++)
    Mask->emplace_back(i);
  std::sort(Mask->begin(), Mask->end(),
                [&UseOrder](unsigned Left, unsigned Right) {
                return UseOrder[Left] < UseOrder[Right];
              });  
}

Instead of having three "if(Mask)"-s.
Sure, you get another for loop, but I think it's cleaner this way.

1098 ↗

(On Diff #89610)

Ok. I think we can do with 1, but I'm still not sure it's an improvement, so we can keep this as is.

lib/Transforms/Vectorize/SLPVectorizer.cpp

2640

Right, that's more or less what I thought.

[• ashahid]

lib/Transforms/Vectorize/SLPVectorizer.cpp

2640

Oh I missed one thing, I think you added a test jumbled-same.ll for a recent fix. I think this test needs to be updated now to account for the use of shuffled value of the loaded vector in "extractelement"?

[mkuper]

lib/Transforms/Vectorize/SLPVectorizer.cpp

2640

Argh, yes, I didn't notice the shuffle is used in the icmp, but not in the extract.
Go ahead and update it. Thanks!

[• ashahid]

Updated the patch to do a proper lane computation for external uses.

[mkuper]

lib/Transforms/Vectorize/SLPVectorizer.cpp

2874

This is the only change in the patch w.r.t the previous version, right?

2877

Isn't this condition equivalent to E->ShuffleMask[i] == ExternalUse.Lane? I think that's a bit clearer, maybe.
(I'm not really certain it's clearer, but I'd like to make sure I understand what's going on here)

Regardless, please add a comment describing what this whole thing does.

test/Transforms/SLPVectorizer/X86/jumbled-same.ll

16 ↗

(On Diff #90429)

Ok, now this makes sense. :-)

[• ashahid]

lib/Transforms/Vectorize/SLPVectorizer.cpp

2874

Yes

2877

Yes, this condition is equivalent to E->ShuffleMask[i] == ExternalUse.Lane and it does look clearer. Oh, sure I will add the comment.

lib/Transforms/Vectorize/SLPVectorizer.cpp

struct TreeEntry {
     }
 
     /// A vector of scalars.
     ValueList Scalars;
 
     /// The Scalars are vectorized into this value. It is initialized to Null.
     Value *VectorizedValue;
 
+    /// A vector of scalars in order uses in program.
+    ValueList InOdrScalars;
+
     /// Do we need to gather this sequence ?
     bool NeedToGather;
 
     /// Do we need to shuffle the load ?
     bool NeedToShuffle;
   };
 
   /// Create a new VectorizableTree entry.
   TreeEntry *newTreeEntry(ArrayRef<Value *> VL, bool Vectorized,
-                          bool NeedToShuffle) {
+                          bool NeedToShuffle,
+                          ArrayRef<Value *> InOdrVL = NULL) {
     VectorizableTree.emplace_back();
     int idx = VectorizableTree.size() - 1;
     TreeEntry *Last = &VectorizableTree[idx];
     Last->Scalars.insert(Last->Scalars.begin(), VL.begin(), VL.end());
     Last->NeedToGather = !Vectorized;
     Last->NeedToShuffle = NeedToShuffle;
+    if (NeedToShuffle)
+      Last->InOdrScalars.insert(Last->InOdrScalars.begin(), InOdrVL.begin(),
+                                InOdrVL.end());
     if (Vectorized) {
       for (int i = 0, e = VL.size(); i != e; ++i) {
         assert(!ScalarToTreeEntry.count(VL[i]) && "Scalar already in tree!");
         ScalarToTreeEntry[VL[i]] = idx;
       }
     } else {
       MustGather.insert(VL.begin(), VL.end());
     }
case Instruction::Load: {
           auto NewVL = makeArrayRef(Sorted.begin(), Sorted.end());
           for (unsigned i = 0, e = NewVL.size() - 1; i < e; ++i) {
             if (!isConsecutiveAccess(NewVL[i], NewVL[i + 1], *DL, *SE)) {
               ShuffledLoads = false;
               break;
             }
           }
           if (ShuffledLoads) {
-            newTreeEntry(NewVL, true, true);
+            newTreeEntry(NewVL, true, true, VL);
             return;
           }
         }
       }
 
       BS.cancelScheduling(VL);
       newTreeEntry(VL, false, false);
 
case Instruction::Load: {
             }
           }
         }
 
         // Generate shuffle for jumbled memory access
         Value *Undef = UndefValue::get(VecTy);
         Value *Shuf = Builder.CreateShuffleVector((Value *)LI, Undef,
                                                   ConstantVector::get(Mask));
+        E->VectorizedValue = Shuf;
+        ++NumVectorInstructions;
         return Shuf;
       }
 
       return LI;
     }
     case Instruction::Store: {
       StoreInst *SI = cast<StoreInst>(VL0);
       unsigned Alignment = SI->getAlignment();

 BoUpSLP::vectorizeTree(MapVector<Value *, DebugLoc> &ExternallyUsedValues) {
 
   // All blocks must be scheduled before any instructions are inserted.
   for (auto &BSIter : BlocksSchedules) {
     scheduleBlock(BSIter.second.get());
   }
 
   Builder.SetInsertPoint(&F->getEntryBlock().front());
-  auto *VectorRoot = vectorizeTree(ArrayRef<Value *>(), &VectorizableTree[0]);
+  auto *VectorRoot = vectorizeTree(VectorizableTree[0].InOdrScalars, &VectorizableTree[0]);
 
   // If the vectorized tree can be rewritten in a smaller type, we truncate the
   // vectorized root. InstCombine will then rewrite the entire expression. We
   // sign extend the extracted values below.
   auto *ScalarRoot = VectorizableTree[0].Scalars[0];
   if (MinBWs.count(ScalarRoot)) {
     if (auto *I = dyn_cast<Instruction>(VectorRoot))
       Builder.SetInsertPoint(&*++BasicBlock::iterator(I));
for (const auto &ExternalUse : ExternalUses) {
     TreeEntry *E = &VectorizableTree[Idx];
     assert(!E->NeedToGather && "Extracting from a gather list");
 
     Value *Vec = E->VectorizedValue;
     assert(Vec && "Can't find vectorizable value");
 
     Value *Lane = Builder.getInt32(ExternalUse.Lane);
     // If User == nullptr, the Scalar is used as extra arg. Generate
     // ExtractElement instruction and update the record for this scalar in
     // ExternallyUsedValues.
     if (!User) {
       assert(ExternallyUsedValues.count(Scalar) &&
              "Scalar with nullptr as an external user must be registered in "
              "ExternallyUsedValues map");
       DebugLoc DL = ExternallyUsedValues[Scalar];
       if (auto *VecI = dyn_cast<Instruction>(Vec)) {
         Builder.SetInsertPoint(VecI->getParent(),
                                std::next(VecI->getIterator()));
       } else {
         Builder.SetInsertPoint(&F->getEntryBlock().front());

test/Transforms/SLPVectorizer/X86/jumbled-load-bug.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -S -slp-vectorizer | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define <4 x i32> @zot() #0 {
+; CHECK-LABEL: @zot(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[P0:%.*]] = getelementptr inbounds [4 x i8], [4 x i8]* undef, i64 undef, i64 0
+; CHECK-NEXT:    [[P1:%.*]] = getelementptr inbounds [4 x i8], [4 x i8]* undef, i64 undef, i64 1
+; CHECK-NEXT:    [[P2:%.*]] = getelementptr inbounds [4 x i8], [4 x i8]* undef, i64 undef, i64 2
+; CHECK-NEXT:    [[P3:%.*]] = getelementptr inbounds [4 x i8], [4 x i8]* undef, i64 undef, i64 3
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i8* [[P0]] to <4 x i8>*
+; CHECK-NEXT:    [[TMP1:%.*]] = load <4 x i8>, <4 x i8>* [[TMP0]], align 1
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <4 x i8> [[TMP1]], <4 x i8> undef, <4 x i32> <i32 1, i32 0, i32 2, i32 3>
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <4 x i8> [[TMP2]], i32 0
+; CHECK-NEXT:    [[I0:%.*]] = insertelement <4 x i8> undef, i8 [[TMP3]], i32 0
+; CHECK-NEXT:    [[TMP4:%.*]] = extractelement <4 x i8> [[TMP2]], i32 1
+; CHECK-NEXT:    [[I1:%.*]] = insertelement <4 x i8> [[I0]], i8 [[TMP4]], i32 1
+; CHECK-NEXT:    [[TMP5:%.*]] = extractelement <4 x i8> [[TMP2]], i32 2
+; CHECK-NEXT:    [[I2:%.*]] = insertelement <4 x i8> [[I1]], i8 [[TMP5]], i32 2
+; CHECK-NEXT:    [[TMP6:%.*]] = extractelement <4 x i8> [[TMP2]], i32 3
+; CHECK-NEXT:    [[I3:%.*]] = insertelement <4 x i8> [[I2]], i8 [[TMP6]], i32 3
+; CHECK-NEXT:    [[RET:%.*]] = zext <4 x i8> [[I3]] to <4 x i32>
+; CHECK-NEXT:    ret <4 x i32> [[RET]]
+;
+bb:
+  %p0 = getelementptr inbounds [4 x i8], [4 x i8]* undef, i64 undef, i64 0
+  %p1 = getelementptr inbounds [4 x i8], [4 x i8]* undef, i64 undef, i64 1
+  %p2 = getelementptr inbounds [4 x i8], [4 x i8]* undef, i64 undef, i64 2
+  %p3 = getelementptr inbounds [4 x i8], [4 x i8]* undef, i64 undef, i64 3
+  %v3 = load i8, i8* %p3, align 1
+  %v2 = load i8, i8* %p2, align 1
+  %v0 = load i8, i8* %p0, align 1
+  %v1 = load i8, i8* %p1, align 1
+  %i0 = insertelement <4 x i8> undef, i8 %v1, i32 0
+  %i1 = insertelement <4 x i8> %i0, i8 %v0, i32 1
+  %i2 = insertelement <4 x i8> %i1, i8 %v2, i32 2
+  %i3 = insertelement <4 x i8> %i2, i8 %v3, i32 3
+  %ret = zext <4 x i8> %i3 to <4 x i32>
+  ret <4 x i32> %ret
+}
+