[LV] Allow assume calls in predicated blocks.

[fhahn]

Summary

The assume intrinsic is intentionally marked as may reading/writing
memory, to avoid passes moving them around. When flattening the CFG
for predicated blocks, we have to drop the assume calls, as they
are control-flow dependent.

Note that this patch also moves setting State->CFG.PrevVPBB after
we executed the recipes in a VPBB. This way, we have access to the real
VPBB while executing a block.

Fixes PR43620.

[xbolva00]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

904

You can drop m_Value()

if (match(&I, m_Intrinsic<Intrinsic::assume>())) {

Wil work too

[hfinkel]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

7004

Unintentional whitespace change?

[fhahn]

Address comments, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

904

Excellent, thanks!

[fhahn]

llvm/test/Transforms/LoopVectorize/predicate-assume.ll

2 ↗

(On Diff #224464)

I've not added a test with calls to assume in the loop header here, because we already have them in llvm/test/Transforms/LoopVectorize/X86/assume.ll.

[Ayal]

If conditional assumes are to be dropped, better do so on entry to VPlan, as in DeadInstructions, rather than representing them in ReplicateRecipe (as do unconditional assumes) and silencing their code generation.

To retain conditional assumes along with their control flow, they could be marked under isScalarWithPredication; but this complicates vectorization, plus what use are such assumes when all else is if-converted(?)

Conditional assumes under uniform control flow could be retained, along with the uniform control flow they depend upon; this may be mostly relevant for outerloop vectorization.

[hsaito]

In D68814#1707477, @Ayal wrote:

If conditional assumes are to be dropped, better do so on entry to VPlan, as in DeadInstructions, rather than representing them in ReplicateRecipe (as do unconditional assumes) and silencing their code generation.

To retain conditional assumes along with their control flow, they could be marked under isScalarWithPredication; but this complicates vectorization, plus what use are such assumes when all else is if-converted(?)

Conditional assumes under uniform control flow could be retained, along with the uniform control flow they depend upon; this may be mostly relevant for outerloop vectorization.

To be more precise:

Assumes should be dropped at the time the vectorizer makes the incompatible decision ---- i.e. a assume is no longer "to be guarded" by the original (per-element or uniform) condition. Right now, in the non-VPlan-native path, recording of the decisions already made happens at entry to VPlan. In the VPlan native path, this would be VPlan to VPlan xform, as the decisions are recorded as they are made.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

155

Do we want to keep this white space?

[fhahn]

In D68814#1707477, @Ayal wrote:

If conditional assumes are to be dropped, better do so on entry to VPlan, as in DeadInstructions, rather than representing them in ReplicateRecipe (as do unconditional assumes) and silencing their code generation.

To retain conditional assumes along with their control flow, they could be marked under isScalarWithPredication; but this complicates vectorization, plus what use are such assumes when all else is if-converted(?)

Conditional assumes under uniform control flow could be retained, along with the uniform control flow they depend upon; this may be mostly relevant for outerloop vectorization.

I thought it might be desirable to only drop the calls at the point where we know for certain that we cannot preserve them: when we emit them in a block that gets merge in its predecessor. By doing it in the recipe, it should also be applicable to the outer loop vectorization. Alternatively we can drop them all as a first step, if that is preferred. Please let me know :)

I think we should try too hard to retain them (through marking them as scalar-with-predication), as the benefit of them for the vectorized loop after the vectorizer is likely to be quite small.

[Ayal]

In D68814#1716053, @fhahn wrote:

In D68814#1707477, @Ayal wrote:

If conditional assumes are to be dropped, better do so on entry to VPlan, as in DeadInstructions, rather than representing them in ReplicateRecipe (as do unconditional assumes) and silencing their code generation.

To retain conditional assumes along with their control flow, they could be marked under isScalarWithPredication; but this complicates vectorization, plus what use are such assumes when all else is if-converted(?)

Conditional assumes under uniform control flow could be retained, along with the uniform control flow they depend upon; this may be mostly relevant for outerloop vectorization.

I thought it might be desirable to only drop the calls at the point where we know for certain that we cannot preserve them: when we emit them in a block that gets merge in its predecessor. By doing it in the recipe, it should also be applicable to the outer loop vectorization. Alternatively we can drop them all as a first step, if that is preferred. Please let me know :)

I vote to simply drop conditional assumes as a first step, at this time. In any case it's better in general if decisions get represented explicitly in VPlan, rather than having such logic kick in during code-gen. (Even if a block gets merged in its predecessor, better also check that the latter has more than one successor.) Following D68577 such assumes could be dropped as a VPlan-to-VPlan transformation.

I think we should try too hard to retain them (through marking them as scalar-with-predication), as the benefit of them for the vectorized loop after the vectorizer is likely to be quite small.

should >> should not?

For the sake of completeness, note that a conditional assume could be retained even under predication, by combining its block predicate within the assume predicate. I.e., a conditional "if (c) assume (p)" can be converted to an unconditional "assume (c implies p)" or "assume (!c || p)".

[fhahn]

Thanks, I’ll update the patch to drop the assume calls for now and let’s go from there

[fhahn]

Updated to use Vplan2VPlan transform to drop assume calls, if required.

[merge_guards_bot]

Build result: FAILURE - Could not check out parent git hash "dcb1398d042ac24cdf8ff0ff7ac7ad2e310ac2dd". It was not found in the repository. Did you configure the "Parent Revision" in Phabricator properly? Trying to apply the patch to the master branch instead...

ERROR: arc patch failed with error code 1. Check build log for details.
Log files: console-log.txt, CMakeCache.txt

[gilr]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

89 ↗

(On Diff #231111)

Shouldn't this recurse into regions for completeness?

98 ↗

(On Diff #231111)

This block is only asserting. Can it be compressed into the asserted expression / go under an #ifdef assert / use llvm_unreachable()?

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

32 ↗

(On Diff #231111)

dropAssumes()?

llvm/test/Transforms/LoopVectorize/predicate-assume.ll

10 ↗

(On Diff #231111)

Enough to CHECK for VPlan and then CHECK-NOT for the assert?

39 ↗

(On Diff #231111)

Enough to CHECK for vector.body and then CHECK-NOT for the assert?

[Ayal]

In D68814#1717282, @Ayal wrote:

[snip]

I vote to simply drop conditional assumes as a first step, at this time. [snip] Following D68577 such assumes could be dropped as a VPlan-to-VPlan transformation.

The alternative mentioned above of following D68577 would result in something like

+ auto &ConditionalAssumes = Legal->getConditionalAssumes();
...
  // Mark instructions we'll need to discard later as
  // ingredients whose recipe we'll need to record.
+  for (auto ConditionalAssume : ConditionalAssumes)
+    RecipeBuilder.recordRecipeOf(ConditionalAssume);
+
...
+  // Discard conditional assumes.
+  for (auto ConditionalAssume : ConditionalAssumes)
+    RecipeBuilder.getRecipe(ConditionalAssume)->eraseFromParent();

where Legal->getConditionalAssumes() is analogous to the proposed Legal->mustDropAssumes(), returning the set of conditional assumes rather than boolean.

Hence the vote to simply include all ConditionalAssumes in DeadInstructions instead.

Having a VPInstruction-based VPlan-to-VPlan transformation would be fine, of course, which drops assumes in if-converted blocks and keeps those under (retained, if any) uniform branches. For buildVPlanWithVPRecipes, however, it would be better to make progress on refactoring Recipes into VPInstructions, wrt assume intrinsics in this case, than to start befriending Recipes and examine their ingredients. As for the VPInstruction-based transformation, should VPInstruction provide an adequate interface (isAssume(), or perhaps an Assume Opcode?) rather than befriending it and examining its UnderlyingInstr(?).

[fhahn]

In D68814#1797963, @Ayal wrote:

[snip]
Hence the vote to simply include all ConditionalAssumes in DeadInstructions instead.

Having a VPInstruction-based VPlan-to-VPlan transformation would be fine, of course, which drops assumes in if-converted blocks and keeps those under (retained, if any) uniform branches. For buildVPlanWithVPRecipes, however, it would be better to make progress on refactoring Recipes into VPInstructions, wrt assume intrinsics in this case, than to start befriending Recipes and examine their ingredients. As for the VPInstruction-based transformation, should VPInstruction provide an adequate interface (isAssume(), or perhaps an Assume Opcode?) rather than befriending it and examining its UnderlyingInstr(?).

Agreed! Let's keep things simple in this patch.

I went ahead and updated the patch to collect the assume instructions in Legal and then add them to the dead instructions just before creating the recipies. Alternatively we could record and remove the assume instructions, but it seems using DeadInstructions is slightly more straight forward.

[merge_guards_bot]

Unit tests: fail. 61241 tests passed, 3 failed and 736 were skipped.

failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_mutex_requirements_mutex/thread_mutex_class/try_lock.pass.cpp
failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_mutex_requirements_mutex/thread_mutex_recursive/lock.pass.cpp
failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_mutex_requirements_mutex/thread_mutex_recursive/try_lock.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[Ayal]

Thanks. Perhaps it's better for Legal to declare via getConditionalAssumes() what is being returning, and let LVP decide what to do with them(*), rather than to assume from the start via getAssumesToDrop() that they must all be dropped.

LVP could collect all conditional assumes itself, btw, but it's easier for Legal to do so, and also logical as they deserve special attention in terms of legality.

(*) As discussed, instead of dropping such assumes could be retained if their condition is retained (being uniform or if handled via isScalarWithPredication()/sinkScalarOperands()?) or if their condition is fused into the assumption.

[fhahn]

In D68814#1804951, @Ayal wrote:

Thanks. Perhaps it's better for Legal to declare via getConditionalAssumes() what is being returning, and let LVP decide what to do with them(*), rather than to assume from the start via getAssumesToDrop() that they must all be dropped.

LVP could collect all conditional assumes itself, btw, but it's easier for Legal to do so, and also logical as they deserve special attention in terms of legality.

+1, getAssumesToDrop makes more sense when we improve things in the future. I've updated the name.

[merge_guards_bot]

Unit tests: pass. 61250 tests passed, 0 failed and 736 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[gilr]

llvm/test/Transforms/LoopVectorize/predicate-assume.ll

10 ↗

(On Diff #231111)

ping?

39 ↗

(On Diff #231111)

ping?

2 ↗

(On Diff #224464)

It would be better to add the new predicated cases to the existing assume.ll.
(Not sure why it resides under X86 - seems it could instead be target independent and force vectorization, but that can be done separately)

[fhahn]

Remove auto-generated check lines.

[fhahn]

In D68814#1804951, @Ayal wrote:

Thanks. Perhaps it's better for Legal to declare via getConditionalAssumes() what is being returning, and let LVP decide what to do with them(*), rather than to assume from the start via getAssumesToDrop() that they must all be dropped.

LVP could collect all conditional assumes itself, btw, but it's easier for Legal to do so, and also logical as they deserve special attention in terms of legality.

(*) As discussed, instead of dropping such assumes could be retained if their condition is retained (being uniform or if handled via isScalarWithPredication()/sinkScalarOperands()?) or if their condition is fused into the assumption.

llvm/test/Transforms/LoopVectorize/predicate-assume.ll

10 ↗

(On Diff #231111)

Ah sorry, I totally missed that. I agree it's better to just check for assumes. The autogenerated checks are unnecessary verbose.

2 ↗

(On Diff #224464)

I am not sure. I personally prefer smaller, more targeted test files, but if you think having a larger test file is better I'll merge them.

[merge_guards_bot]

Unit tests: pass. 61257 tests passed, 0 failed and 736 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[Ayal]

Agree that it's better to merge the new conditional assume tests with the existing unconditional assume tests in the same file.

Thanks for making the changes!

llvm/test/Transforms/LoopVectorize/predicate-assume.ll

2 ↗

(On Diff #224464)

The above forthcoming comment emphasizes that these tests complement each other: “I've not added a test with calls to assume in the loop header here, because we already have them in llvm/test/Transforms/LoopVectorize/X86/assume.ll”.

Quoting the end of https://llvm.org/docs/TestingGuide.html#writing-new-regression-tests:
“Put related tests into a single file rather than having a separate file per test. Check if there are files already covering your feature and consider adding your code there instead of creating a new file.”

[fhahn]

In D68814#1815837, @Ayal wrote:

Agree that it's better to merge the new conditional assume tests with the existing unconditional assume tests in the same file.

Sure! I've made the X86/assume.ll test non-X86 specific and I'll land the patch with the new test added to it directly.

Thanks for taking a look!

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

public:
   bool isMaskRequired(const Instruction *I) { return (MaskedOp.count(I) != 0); }
 
   unsigned getNumStores() const { return LAI->getNumStores(); }
   unsigned getNumLoads() const { return LAI->getNumLoads(); }
 
   // Returns true if the NoNaN attribute is set on the function.
   bool hasFunNoNaNAttr() const { return HasFunNoNaNAttr; }
 
+  /// Returns all assume calls in predicated blocks. They need to be dropped
+  /// when flattening the CFG.
+  const SmallPtrSetImpl<Instruction *> &getConditionalAssumes() const {
+    return ConditionalAssumes;
+  }
+
 private:
   /// Return true if the pre-header, exiting and latch blocks of \p Lp and all
   /// its nested loops are considered legal for vectorization. These legal
   /// checks are common for inner and outer loop vectorization.
   /// Temporarily taking UseVPlanNativePath parameter. If true, take
   /// the new code path being implemented for outer loop vectorization
   /// (should be functional for inner loop vectorization) based on VPlan.
   /// If false, good old LV code.
private:
 
   /// The assumption cache analysis is used to compute the minimum type size in
   /// which a reduction can be computed.
   AssumptionCache *AC;
 
   /// While vectorizing these instructions we have to generate a
   /// call to the appropriate masked intrinsic
   SmallPtrSet<const Instruction *, 8> MaskedOp;
+
+  /// Assume instructions in predicated blocks must be dropped if the CFG gets
+  /// flattened.
+  SmallPtrSet<Instruction *, 8> ConditionalAssumes;
 };
 
 } // namespace llvm
 
 #endif // LLVM_TRANSFORMS_VECTORIZE_LOOPVECTORIZATIONLEGALITY_H

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

 //===- LoopVectorizationLegality.cpp --------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file provides loop vectorization legality analysis. Original code
 // resided in LoopVectorize.cpp for a long time.
 //
 // At this point, it is implemented as a utility class, not as an analysis
 // pass. It should be easy to create an analysis pass around it if there
 // is a need (but D45420 needs to happen first).
 //
-#include "llvm/Transforms/Vectorize/LoopVectorize.h"
 #include "llvm/Transforms/Vectorize/LoopVectorizationLegality.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/PatternMatch.h"
+#include "llvm/Transforms/Vectorize/LoopVectorize.h"
 
 using namespace llvm;
+using namespace PatternMatch;
 
 #define LV_NAME "loop-vectorize"
 #define DEBUG_TYPE LV_NAME
 
 extern cl::opt<bool> EnableVPlanPredication;
 
 static cl::opt<bool>
     EnableIfConversion("enable-if-conversion", cl::init(true), cl::Hidden,
bool LoopVectorizationLegality::blockCanBePredicated(
 
   for (Instruction &I : *BB) {
     // Check that we don't have a constant expression that can trap as operand.
     for (Value *Operand : I.operands()) {
       if (auto *C = dyn_cast<Constant>(Operand))
         if (C->canTrap())
           return false;
     }
+
+    // We can predicate blocks with calls to assume, as long as we drop them in
+    // case we flatten the CFG via predication.
+    if (match(&I, m_Intrinsic<Intrinsic::assume>())) {
+      ConditionalAssumes.insert(&I);
+      continue;
+    }
+
     // We might be able to hoist the load.
     if (I.mayReadFromMemory()) {
       auto *LI = dyn_cast<LoadInst>(&I);
       if (!LI)
         return false;
       if (!SafePtrs.count(LI->getPointerOperand())) {
         // !llvm.mem.parallel_loop_access implies if-conversion safety.
         // Otherwise, record that the load needs (real or emulated) masking

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

 #include <iterator>
 #include <limits>
 #include <memory>
 #include <string>
 #include <tuple>
 #include <utility>
 
 using namespace llvm;
 
 #define LV_NAME "loop-vectorize"
 #define DEBUG_TYPE LV_NAME
 
 /// @{
 /// Metadata attribute names
 static const char *const LLVMLoopVectorizeFollowupAll =
     "llvm.loop.vectorize.followup_all";
 static const char *const LLVMLoopVectorizeFollowupVectorized =
if (auto *PredInst = dyn_cast<Instruction>(Op))
         PredInst2Recipe[PredInst]->setAlsoPack(false);
 
   // Finalize the recipe for Instr, first if it is not predicated.
   if (!IsPredicated) {
     LLVM_DEBUG(dbgs() << "LV: Scalarizing:" << *I << "\n");
     VPBB->appendRecipe(Recipe);
     return VPBB;
   }
   LLVM_DEBUG(dbgs() << "LV: Scalarizing and predicating:" << *I << "\n");
   assert(VPBB->getSuccessors().empty() &&
          "VPBB has successors when handling predicated replication.");
   // Record predicated instructions for above packing optimizations.
   PredInst2Recipe[I] = Recipe;
   VPBlockBase *Region = createReplicateRegion(I, Recipe, Plan);
   VPBlockUtils::insertBlockAfter(Region, VPBB);
   auto *RegSucc = new VPBasicBlock();
   VPBlockUtils::insertBlockAfter(RegSucc, Region);
for (unsigned i = 0; i < IG->getFactor(); i++)
         RecipeBuilder.recordRecipeOf(Member);
   };
 
   // ---------------------------------------------------------------------------
   // Build initial VPlan: Scan the body of the loop in a topological order to
   // visit each basic block after having visited its predecessor basic blocks.
   // ---------------------------------------------------------------------------
 
+  // Add assume instructions we need to drop to DeadInstructions, to prevent
+  // them from being added to the VPlan.
+  // TODO: We only need to drop assumes in blocks that get flattend. If the
+  // control flow is preserved, we should keep them.
+  auto &ConditionalAssumes = Legal->getConditionalAssumes();
+  DeadInstructions.insert(ConditionalAssumes.begin(), ConditionalAssumes.end());
+
   // Create a dummy pre-entry VPBasicBlock to start building the VPlan.
   VPBasicBlock *VPBB = new VPBasicBlock("Pre-Entry");
   auto Plan = std::make_unique<VPlan>(VPBB);
 
   // Represent values that will have defs inside VPlan.
   for (Value *V : NeedDef)
     Plan->addVPValue(V);
 

llvm/test/Transforms/LoopVectorize/assume.ll

for.body:                                         ; preds = %for.body, %entry
   store float %add, float* %arrayidx5, align 4
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %exitcond = icmp eq i64 %indvars.iv, 1599
   br i1 %exitcond, label %for.end, label %for.body
 
 for.end:                                          ; preds = %for.body
   ret void
 }
+
+; Test case for PR43620. Make sure we can vectorize with predication in presence
+; of assume calls. For now, check that we drop all assumes in predicated blocks
+; in the vector body.
+define void @predicated_assume(float* noalias nocapture readonly %a, float* noalias nocapture %b, i32 %n) {
+
+; Check that the vector.body does not contain any assumes.
+
+; CHECK-LABEL: @predicated_assume(
+; CHECK:     vector.body:
+; CHECK-NOT:   call void @llvm.assume(
+; CHECK:     scalar.ph:
+;
+entry:
+  %cmp15 = icmp eq i32 %n, 0
+  br i1 %cmp15, label %for.cond.cleanup, label %for.body.preheader
+
+for.body.preheader:                               ; preds = %entry
+  %0 = zext i32 %n to i64
+  br label %for.body
+
+for.cond.cleanup.loopexit:                        ; preds = %if.end5
+  br label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.cond.cleanup.loopexit, %entry
+  ret void
+
+for.body:                                         ; preds = %for.body.preheader, %if.end5
+  %indvars.iv = phi i64 [ 0, %for.body.preheader ], [ %indvars.iv.next, %if.end5 ]
+  %cmp1 = icmp ult i64 %indvars.iv, 495616
+  br i1 %cmp1, label %if.end5, label %if.else
+
+if.else:                                          ; preds = %for.body
+  %cmp2 = icmp ult i64 %indvars.iv, 991232
+  tail call void @llvm.assume(i1 %cmp2)
+  br label %if.end5
+
+if.end5:                                          ; preds = %for.body, %if.else
+  %x.0 = phi float [ 4.200000e+01, %if.else ], [ 2.300000e+01, %for.body ]
+  %arrayidx = getelementptr inbounds float, float* %a, i64 %indvars.iv
+  %1 = load float, float* %arrayidx, align 4
+  %mul = fmul float %x.0, %1
+  %arrayidx7 = getelementptr inbounds float, float* %b, i64 %indvars.iv
+  store float %mul, float* %arrayidx7, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %cmp = icmp eq i64 %indvars.iv.next, %0
+  br i1 %cmp, label %for.cond.cleanup.loopexit, label %for.body
+}