[LoopVectorizer] Support predicating instructions in blocks with several input edges

[mkuper]

Summary

It's not really clear to me (or to Gil, who suggested the fix) why this was restricted to single-entry blocks to begin with.

This fixes PR30172.

[mssimpso]

This looks reasonable to me, but I have a question about the test case (see the inline comment). I too cannot say why the original conditional stores patch limited this to single-predecessor blocks.

test/Transforms/LoopVectorize/if-pred-non-void.ll

197

Is the test checking what you want it to? I was expecting if.then to have more than one predecessor (and trigger the assert) since it's the block with the sdiv's.

[mkuper]

test/Transforms/LoopVectorize/if-pred-non-void.ll

197

You're completely right - I was tweaking the test to make the checks more readable, and didn't notice I no longer had the right structure.
Will fix it, thanks!

[gilr]

test/Transforms/LoopVectorize/if-pred-non-void.ll

195

It might make the test a bit more future-proof to base the branch on a non loop-invariant condition - make it more resistant to versioning/unswitching.

[mkuper]

Now, with a correct test.

[mssimpso]

LGTM!

[gilr]

LGTM - Thanks, Michael!

lib/Transforms/Vectorize/LoopVectorize.cpp

void InnerLoopVectorizer::scalarizeInstruction(Instruction *Instr,
 
   // Does this instruction return a value ?
   bool IsVoidRetTy = Instr->getType()->isVoidTy();
 
   // Initialize a new scalar map entry.
   ScalarParts Entry(UF);
 
   VectorParts Cond;
-  if (IfPredicateInstr) {
-    assert(Instr->getParent()->getSinglePredecessor() &&
-           "Only support single predecessor blocks");
-    Cond = createEdgeMask(Instr->getParent()->getSinglePredecessor(),
-                          Instr->getParent());
-  }
+  if (IfPredicateInstr)
+    Cond = createBlockInMask(Instr->getParent());
 
   // For each vector unroll 'part':
   for (unsigned Part = 0; Part < UF; ++Part) {
     Entry[Part].resize(VF);
     // For each scalar that we create:
     for (unsigned Width = 0; Width < VF; ++Width) {
 
       // Start if-block.
void InnerLoopUnroller::scalarizeInstruction(Instruction *Instr,
 
   // Does this instruction return a value ?
   bool IsVoidRetTy = Instr->getType()->isVoidTy();
 
   // Initialize a new scalar map entry.
   ScalarParts Entry(UF);
 
   VectorParts Cond;
-  if (IfPredicateInstr) {
-    assert(Instr->getParent()->getSinglePredecessor() &&
-           "Only support single predecessor blocks");
-    Cond = createEdgeMask(Instr->getParent()->getSinglePredecessor(),
-                          Instr->getParent());
-  }
+  if (IfPredicateInstr)
+    Cond = createBlockInMask(Instr->getParent());
 
   // For each vector unroll 'part':
   for (unsigned Part = 0; Part < UF; ++Part) {
     Entry[Part].resize(1);
     // For each scalar that we create:
 
     // Start an "if (pred) a[i] = ..." block.
     Value *Cmp = nullptr;

test/Transforms/LoopVectorize/if-pred-non-void.ll

 
 if.end:                                           ; preds = %if.then, %for.body
   %ysd.0 = phi i32 [ %rsd, %if.then ], [ %psd, %for.body ]
   store i32 %ysd.0, i32* %isd, align 4
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %exitcond = icmp eq i64 %indvars.iv.next, 128
   br i1 %exitcond, label %for.cond.cleanup, label %for.body
 }
+
+define void @pr30172(i32* nocapture %asd, i32* nocapture %bsd, i1 %p) {
+entry:
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %if.end
+  ret void
+
+; CHECK-LABEL: pr30172
+; CHECK: vector.ph:
+; CHECK: %[[SI:.+]] = insertelement <2 x i1> undef, i1 %p, i32 0
+; CHECK: %[[P:.+]] = shufflevector <2 x i1> %[[SI]], <2 x i1> undef, <2 x i32> zeroinitializer
+; CHECK: vector.body:
+; CHECK: %[[CMP:.+]] = icmp slt <2 x i32> %{{.*}}, <i32 100, i32 100>
+; CHECK: %[[AND1:.+]] = and <2 x i1> %[[CMP]], <i1 true, i1 true>
+; CHECK: %[[OR1:.+]] = or <2 x i1> zeroinitializer, %[[AND1]]
+; CHECK: %[[AND2:.+]] = and <2 x i1> %[[P]], %[[OR1]]
+; CHECK: %[[OR2:.+]] = or <2 x i1> zeroinitializer, %[[AND2]]
+; CHECK: %[[EXTRACT:.+]] = extractelement <2 x i1> %[[OR2]], i32 0
+; CHECK: %[[MASK:.+]] = icmp eq i1 %[[EXTRACT]], true
+; CHECK: br i1 %[[MASK]], label %[[THEN:[a-zA-Z0-9.]+]], label %[[FI:[a-zA-Z0-9.]+]]
+; CHECK: [[THEN]]:
+; CHECK:   %[[PD:[a-zA-Z0-9]+]] = sdiv i32 %{{.*}}, %{{.*}}
+; CHECK:   br label %[[FI]]
+; CHECK: [[FI]]:
+; CHECK:   %{{.*}} = phi i32 [ undef, %vector.body ], [ %[[PD]], %[[THEN]] ]
+
+
+for.body:                                         ; preds = %if.end, %entry
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %if.end ]
+  %isd = getelementptr inbounds i32, i32* %asd, i64 %indvars.iv
+  %lsd = load i32, i32* %isd, align 4
+  %isd.b = getelementptr inbounds i32, i32* %bsd, i64 %indvars.iv
+  %lsd.b = load i32, i32* %isd.b, align 4
+  %psd = add nsw i32 %lsd, 23
+  %cmp1 = icmp slt i32 %lsd, 100
+  br i1 %cmp1, label %diamond, label %if.end
+
+diamond:
+  br i1 %p, label %if.then, label %if.end
+
+if.then:                                          ; preds = %for.body
+  %sd1 = sdiv i32 %psd, %lsd
+  %rsd = sdiv i32 %lsd.b, %sd1
+  br label %if.end
+
+if.end:                                           ; preds = %if.then, %for.body
+  %ysd.0 = phi i32 [ %rsd, %if.then ], [ %psd, %for.body ], [ %psd, %diamond ] 
+  store i32 %ysd.0, i32* %isd, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond = icmp eq i64 %indvars.iv.next, 128
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body
+}

test/Transforms/LoopVectorize/if-pred-stores.ll

 ; VEC-LABEL: test
 ; VEC:   %[[v0:.+]] = add i64 %index, 0
 ; VEC:   %[[v1:.+]] = add i64 %index, 1
 ; VEC:   %[[v2:.+]] = getelementptr inbounds i32, i32* %f, i64 %[[v0]]
 ; VEC:   %[[v4:.+]] = getelementptr inbounds i32, i32* %f, i64 %[[v1]]
 ; VEC:   %[[v8:.+]] = icmp sgt <2 x i32> %{{.*}}, <i32 100, i32 100>
 ; VEC:   %[[v9:.+]] = add nsw <2 x i32> %{{.*}}, <i32 20, i32 20>
 ; VEC:   %[[v10:.+]] = and <2 x i1> %[[v8]], <i1 true, i1 true>
-; VEC:   %[[v11:.+]] = extractelement <2 x i1> %[[v10]], i32 0
+; VEC:   %[[o1:.+]] = or <2 x i1> zeroinitializer, %[[v10]]
+; VEC:   %[[v11:.+]] = extractelement <2 x i1> %[[o1]], i32 0
 ; VEC:   %[[v12:.+]] = icmp eq i1 %[[v11]], true
 ; VEC:   br i1 %[[v12]], label %[[cond:.+]], label %[[else:.+]]
 ;
 ; VEC: [[cond]]:
 ; VEC:   %[[v13:.+]] = extractelement <2 x i32> %[[v9]], i32 0
 ; VEC:   store i32 %[[v13]], i32* %[[v2]], align 4
 ; VEC:   br label %[[else:.+]]
 ;
 ; VEC: [[else]]:
-; VEC:   %[[v15:.+]] = extractelement <2 x i1> %[[v10]], i32 1
+; VEC:   %[[v15:.+]] = extractelement <2 x i1> %[[o1]], i32 1
 ; VEC:   %[[v16:.+]] = icmp eq i1 %[[v15]], true
 ; VEC:   br i1 %[[v16]], label %[[cond2:.+]], label %[[else2:.+]]
 ;
 ; VEC: [[cond2]]:
 ; VEC:   %[[v17:.+]] = extractelement <2 x i32> %[[v9]], i32 1
 ; VEC:   store i32 %[[v17]], i32* %[[v4]], align 4
 ; VEC:   br label %[[else2:.+]]
 ;
 ; VEC: [[else2]]:
 
 ; UNROLL-LABEL: test
 ; UNROLL: vector.body:
 ; UNROLL:   %[[IND:[a-zA-Z0-9]+]] = add i64 %{{.*}}, 0
 ; UNROLL:   %[[IND1:[a-zA-Z0-9]+]] = add i64 %{{.*}}, 1
 ; UNROLL:   %[[v0:[a-zA-Z0-9]+]] = getelementptr inbounds i32, i32* %f, i64 %[[IND]]
 ; UNROLL:   %[[v1:[a-zA-Z0-9]+]] = getelementptr inbounds i32, i32* %f, i64 %[[IND1]]
 ; UNROLL:   %[[v2:[a-zA-Z0-9]+]] = load i32, i32* %[[v0]], align 4
 ; UNROLL:   %[[v3:[a-zA-Z0-9]+]] = load i32, i32* %[[v1]], align 4
 ; UNROLL:   %[[v4:[a-zA-Z0-9]+]] = icmp sgt i32 %[[v2]], 100
 ; UNROLL:   %[[v5:[a-zA-Z0-9]+]] = icmp sgt i32 %[[v3]], 100
 ; UNROLL:   %[[v6:[a-zA-Z0-9]+]] = add nsw i32 %[[v2]], 20
 ; UNROLL:   %[[v7:[a-zA-Z0-9]+]] = add nsw i32 %[[v3]], 20
-; UNROLL:   %[[v8:[a-zA-Z0-9]+]] = icmp eq i1 %[[v4]], true
+; UNROLL:   %[[o1:[a-zA-Z0-9]+]] = or i1 false, %[[v4]]
+; UNROLL:   %[[o2:[a-zA-Z0-9]+]] = or i1 false, %[[v5]]
+; UNROLL:   %[[v8:[a-zA-Z0-9]+]] = icmp eq i1 %[[o1]], true
 ; UNROLL:   br i1 %[[v8]], label %[[cond:[a-zA-Z0-9.]+]], label %[[else:[a-zA-Z0-9.]+]]
 ;
 ; UNROLL: [[cond]]:
 ; UNROLL:   store i32 %[[v6]], i32* %[[v0]], align 4
 ; UNROLL:   br label %[[else]]
 ;
 ; UNROLL: [[else]]:
-; UNROLL:   %[[v9:[a-zA-Z0-9]+]] = icmp eq i1 %[[v5]], true
+; UNROLL:   %[[v9:[a-zA-Z0-9]+]] = icmp eq i1 %[[o2]], true
 ; UNROLL:   br i1 %[[v9]], label %[[cond2:[a-zA-Z0-9.]+]], label %[[else2:[a-zA-Z0-9.]+]]
 ;
 ; UNROLL: [[cond2]]:
 ; UNROLL:   store i32 %[[v7]], i32* %[[v1]], align 4
 ; UNROLL:   br label %[[else2]]
 ;
 ; UNROLL: [[else2]]: