[SLP]Fix PR55734: SLP vectorizer's reduce_and formation introduces poison.

[ABataev]

Summary

Need either follow the original order of the operands for bool logical
ops, or emit freeze instruction to avoid poison propagation.

[ABataev]

Rebase and use isGuaranteedNotToBePoison as an extra check

[ABataev]

Rebase

[ABataev]

Ping!

[spatel]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

11096

This could use a comment to describe the patterns/logic.

IIUC, we want to create a bool logic op like:

select i1 LHS, i1 RHS, i1 false ; logical and

If we have original code like this:

RedOp1 = select i1 ?, i1 LHS, i1 false
RedOp2 = select i1 RHS, i1 ?, i1 false

Then, we swap LHS/RHS to create a new op that matches the poison semantics of the original code.

If we have original code like this and both values could be poison:

RedOp1 = select i1 ?, i1 LHS, i1 false
RedOp2 = select i1 ?, i1 RHS, i1 false

Then, we must freeze LHS in the new op. Can we manufacture a test for this 2nd path? I don't think there's any coverage for that path in the current test diffs.

[ABataev]

Reworked addition of the reduction. For logical ops, reduction result is added as a first operand,
it is safe to use it as an LHS in scalar logical ands/ors.
For others, it is added as a last operand.

[spatel]

llvm/test/Transforms/SLPVectorizer/X86/reduction-logical.ll

312–313

The updated patch does not fix the bug?
https://alive2.llvm.org/ce/z/YYaTVw

[ABataev]

Fix poison leak for vectorized tree top values.

[ABataev]

Rebase

[ABataev]

Rebase

[ABataev]

Ping!

[spatel]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

10312–10315

I think we should just fix this to include isa<SelectInst>, so it matches the code comment. Then, we don't need the confusing extra variable in this patch. All of the existing uses are checking for SelectInst, so that can be a preliminary cleanup patch.

[ABataev]

Address comment

[spatel]

LGTM

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

class HorizontalReduction {
   /// The type of reduction operation.
   RecurKind RdxKind;
 
   static bool isCmpSelMinMax(Instruction *I) {
     return match(I, m_Select(m_Cmp(), m_Value(), m_Value())) &&
            RecurrenceDescriptor::isMinMaxRecurrenceKind(getRdxKind(I));
   }
 
   // And/or are potentially poison-safe logical patterns like:
   // select x, y, false
   // select x, true, y
   static bool isBoolLogicOp(Instruction *I) {
     return match(I, m_LogicalAnd(m_Value(), m_Value())) ||
            match(I, m_LogicalOr(m_Value(), m_Value()));
   }
 
   /// Checks if instruction is associative and can be vectorized.
   static bool isVectorizable(RecurKind Kind, Instruction *I) {
     if (Kind == RecurKind::None)
       return false;
for (unsigned I = 0, E = ReducedVals.size(); I < E; ++I) {
           ++VectorizedVals.try_emplace(TrackedToOrig.find(V)->second, 0)
                 .first->getSecond();
         Pos += ReduxWidth;
         Start = Pos;
         ReduxWidth = PowerOf2Floor(NumReducedVals - Pos);
       }
     }
     if (VectorizedTree) {
+      // Reorder operands of bool logical op in the natural order to avoid
+      // possible problem with poison propagation. If not possible to reorder
+      // (both operands are originally RHS), emit an extra freeze instruction
+      // for the LHS operand.
+      auto &&FixBoolLogicalOps = [&Builder](Value *&LHS, Value *&RHS,
+                                            Instruction *RedOp1,
+                                            Instruction *RedOp2) {
+        if (!isa<SelectInst>(RedOp1) || !isBoolLogicOp(RedOp1))
+          return;
+        if (getRdxOperand(RedOp1, 0) == LHS ||
+            getRdxOperand(RedOp1, 1) != LHS || isGuaranteedNotToBePoison(LHS))
+          return;
+        if (!isa<SelectInst>(RedOp2) || !isBoolLogicOp(RedOp2))
+          return;
+        if (getRdxOperand(RedOp2, 0) == RHS ||
+            getRdxOperand(RedOp2, 1) != RHS || isGuaranteedNotToBePoison(RHS)) {
+          std::swap(LHS, RHS);
+          return;
+        }
+        LHS = Builder.CreateFreeze(LHS);
+      };
       // Finish the reduction.
       // Need to add extra arguments and not vectorized possible reduction
       // values.
       // Try to avoid dependencies between the scalar remainders after
       // reductions.
       auto &&FinalGen =
-          [this, &Builder,
-           &TrackedVals](ArrayRef<std::pair<Instruction *, Value *>> InstVals) {
+          [this, &Builder, &TrackedVals, &FixBoolLogicalOps](
+              ArrayRef<std::pair<Instruction *, Value *>> InstVals) {
             unsigned Sz = InstVals.size();
             SmallVector<std::pair<Instruction *, Value *>> ExtraReds(Sz / 2 +
                                                                      Sz % 2);
             for (unsigned I = 0, E = (Sz / 2) * 2; I < E; I += 2) {
               Instruction *RedOp = InstVals[I + 1].first;
               Builder.SetCurrentDebugLocation(RedOp->getDebugLoc());
               Value *RdxVal1 = InstVals[I].second;
               Value *StableRdxVal1 = RdxVal1;
               auto It1 = TrackedVals.find(RdxVal1);
               if (It1 != TrackedVals.end())
                 StableRdxVal1 = It1->second;
               Value *RdxVal2 = InstVals[I + 1].second;
               Value *StableRdxVal2 = RdxVal2;
               auto It2 = TrackedVals.find(RdxVal2);
               if (It2 != TrackedVals.end())
                 StableRdxVal2 = It2->second;
+              // To prevent poison from leaking across what used to be
+              // sequential, safe, scalar boolean logic operations, the
+              // reduction operand must be frozen.
+              FixBoolLogicalOps(StableRdxVal1, StableRdxVal2, InstVals[I].first,
+                                RedOp);
               Value *ExtraRed = createOp(Builder, RdxKind, StableRdxVal1,
                                          StableRdxVal2, "op.rdx", ReductionOps);
               ExtraReds[I / 2] = std::make_pair(InstVals[I].first, ExtraRed);
             }
             if (Sz % 2 == 1)
               ExtraReds[Sz / 2] = InstVals.back();
             return ExtraReds;
           };

llvm/test/Transforms/SLPVectorizer/X86/reduction-logical.ll

 ; CHECK-NEXT:    [[TMP4:%.*]] = extractelement <4 x i1> [[TMP1]], i32 1
 ; CHECK-NEXT:    [[S1:%.*]] = select i1 [[TMP3]], i1 [[TMP4]], i1 false
 ; CHECK-NEXT:    [[TMP5:%.*]] = extractelement <4 x i1> [[TMP1]], i32 2
 ; CHECK-NEXT:    [[S2:%.*]] = select i1 [[S1]], i1 [[TMP5]], i1 false
 ; CHECK-NEXT:    call void @use1(i1 [[S2]])
 ; CHECK-NEXT:    [[TMP6:%.*]] = freeze <4 x i1> [[TMP2]]
 ; CHECK-NEXT:    [[TMP7:%.*]] = call i1 @llvm.vector.reduce.and.v4i1(<4 x i1> [[TMP6]])
 ; CHECK-NEXT:    [[TMP8:%.*]] = extractelement <4 x i1> [[TMP1]], i32 3
-; CHECK-NEXT:    [[OP_RDX:%.*]] = select i1 [[TMP8]], i1 [[S2]], i1 false
+; CHECK-NEXT:    [[OP_RDX:%.*]] = select i1 [[S2]], i1 [[TMP8]], i1 false
 ; CHECK-NEXT:    [[OP_RDX1:%.*]] = select i1 [[TMP7]], i1 [[OP_RDX]], i1 false
 ; CHECK-NEXT:    ret i1 [[OP_RDX1]]
 ;
   %x0 = extractelement <4 x i32> %x, i32 0
   %x1 = extractelement <4 x i32> %x, i32 1
   %x2 = extractelement <4 x i32> %x, i32 2
   %x3 = extractelement <4 x i32> %x, i32 3
   %c0 = icmp slt i32 %x0, 42
   %c1 = icmp slt i32 %x1, 42
   %c2 = icmp slt i32 %x2, 42

 ; SSE-NEXT:    [[TMP5:%.*]] = insertelement <2 x i32> [[TMP4]], i32 [[TMP3]], i32 1
 ; SSE-NEXT:    [[TMP6:%.*]] = icmp slt <2 x i32> [[TMP5]], <i32 42, i32 42>
 ; SSE-NEXT:    [[C2:%.*]] = icmp slt i32 [[TMP1]], 42
 ; SSE-NEXT:    [[TMP7:%.*]] = icmp sgt <4 x i32> [[X]], <i32 17, i32 17, i32 17, i32 17>
 ; SSE-NEXT:    [[TMP8:%.*]] = freeze <4 x i1> [[TMP7]]
 ; SSE-NEXT:    [[TMP9:%.*]] = call i1 @llvm.vector.reduce.and.v4i1(<4 x i1> [[TMP8]])
 ; SSE-NEXT:    [[TMP10:%.*]] = extractelement <2 x i1> [[TMP6]], i32 0
 ; SSE-NEXT:    [[TMP11:%.*]] = extractelement <2 x i1> [[TMP6]], i32 1
-; SSE-NEXT:    [[OP_RDX:%.*]] = select i1 [[TMP10]], i1 [[TMP11]], i1 false
+; SSE-NEXT:    [[OP_RDX:%.*]] = select i1 [[TMP11]], i1 [[TMP10]], i1 false
 ; SSE-NEXT:    [[OP_RDX1:%.*]] = select i1 [[OP_RDX]], i1 [[C2]], i1 false
 ; SSE-NEXT:    [[OP_RDX2:%.*]] = select i1 [[TMP9]], i1 [[OP_RDX1]], i1 false
 ; SSE-NEXT:    ret i1 [[OP_RDX2]]
 ;
 ; AVX-LABEL: @logical_and_icmp_clamp_partial(
 ; AVX-NEXT:    [[TMP1:%.*]] = extractelement <4 x i32> [[X:%.*]], i32 2
 ; AVX-NEXT:    [[TMP2:%.*]] = extractelement <4 x i32> [[X]], i32 1
 ; AVX-NEXT:    [[TMP3:%.*]] = extractelement <4 x i32> [[X]], i32 0
 ; AVX-NEXT:    [[C0:%.*]] = icmp slt i32 [[TMP3]], 42
 ; AVX-NEXT:    [[C1:%.*]] = icmp slt i32 [[TMP2]], 42
 ; AVX-NEXT:    [[C2:%.*]] = icmp slt i32 [[TMP1]], 42
 ; AVX-NEXT:    [[TMP4:%.*]] = icmp sgt <4 x i32> [[X]], <i32 17, i32 17, i32 17, i32 17>
 ; AVX-NEXT:    [[TMP5:%.*]] = freeze <4 x i1> [[TMP4]]
 ; AVX-NEXT:    [[TMP6:%.*]] = call i1 @llvm.vector.reduce.and.v4i1(<4 x i1> [[TMP5]])
-; AVX-NEXT:    [[OP_RDX:%.*]] = select i1 [[C1]], i1 [[C0]], i1 false
+; AVX-NEXT:    [[OP_RDX:%.*]] = select i1 [[C0]], i1 [[C1]], i1 false
 ; AVX-NEXT:    [[OP_RDX1:%.*]] = select i1 [[OP_RDX]], i1 [[C2]], i1 false
 ; AVX-NEXT:    [[OP_RDX2:%.*]] = select i1 [[TMP6]], i1 [[OP_RDX1]], i1 false
 ; AVX-NEXT:    ret i1 [[OP_RDX2]]
 ;
   %x0 = extractelement <4 x i32> %x, i32 0
   %x1 = extractelement <4 x i32> %x, i32 1
   %x2 = extractelement <4 x i32> %x, i32 2
   %x3 = extractelement <4 x i32> %x, i32 3