[GVN] Refactor handling of pointer-select in GVN pass

[kachkov98]

Summary

This patch extends Def memory dependency with support of select instructions to consistently handle pointer-select conversion.

In general it would be great to support address translation through select inst (extend PHITransAddr?) to optimize different cases when array index depends from select like: https://godbolt.org/z/TsarK6xEs (I've reported this issue in https://github.com/llvm/llvm-project/issues/58569), now we only handle select-between-pointer; this patch is preparation for this work.

[nikic]

Could you please split out the AnalyzeLoadAvailability() refactoring to std::optional into a separate NFC patch (and base this one on top)?

[kachkov98]

Split AnalyzeLoadAvailability refactoring into separate patch: https://reviews.llvm.org/D141664

[nikic]

Thanks! Could you please also split off the change to findDominatingLoad()? For that change, we also need to introduce a limit on the number of instructions that get walked while finding the dominating value (100 would be consistent with MemDep), otherwise there may be issues with very large blocks. There should also be a test that shows we can now handle a case that we previously couldn't (with the dominating load being in a preceding basic block).

[kachkov98]

Split findDominatingLoad() rework into separate patch

[nikic]

LGTM

It's a bit odd to have a non-memory instruction as the dependency, but this does seem to work out nicely in practice.

llvm/lib/Transforms/Scalar/GVN.cpp

1124

Outdated comment (End)

1134

Outdated comment (End).

[nikic]

It looks like this breaks the build on 32-bit targets, because we don't have enough alignment bits.

[kachkov98]

Since extending MemDepResult looks not feasable (and replacing it to tag-pointer pair also looks bad because it frequently passed by value), I think we can treat select dependency as Def. It also fits naturally into existing GVN infra, and the meaning of Def dependency is already overloaded (it can be other load/store/alloca/memory intrinsic/etc.)

[kachkov98]

Treat select as Def dependency

[nikic]

LGTM

[nikic]

(Patch description needs an update though)

llvm/include/llvm/Analysis/MemoryDependenceAnalysis.h

enum DepType {
     ///      this case, the load is loading an undef value or a store is the
     ///      first store to (that part of) the allocation.
     ///   3. Dependence queries on calls return Def only when they are readonly
     ///      calls or memory use intrinsics with identical callees and no
     ///      intervening clobbers.  No validation is done that the operands to
     ///      the calls are the same.
     Def,
 
+    /// This marker indicates that the query has dependency from select
+    /// instruction.
+    Select,
+
     /// This marker indicates that the query has no known dependency in the
     /// specified block.
     ///
     /// More detailed state info is encoded in the upper part of the pair (i.e.
     /// the Instruction*)
     Other
   };
 
enum OtherType {
     /// This marker indicates that the query dependency is unknown.
     Unknown
   };
 
   using ValueTy = PointerSumType<
       DepType, PointerSumTypeMember<Invalid, Instruction *>,
       PointerSumTypeMember<Clobber, Instruction *>,
       PointerSumTypeMember<Def, Instruction *>,
+      PointerSumTypeMember<Select, Instruction *>,
       PointerSumTypeMember<Other, PointerEmbeddedInt<OtherType, 3>>>;
   ValueTy Value;
 
   explicit MemDepResult(ValueTy V) : Value(V) {}
 
 public:
   MemDepResult() = default;
 
   /// get methods: These are static ctor methods for creating various
   /// MemDepResult kinds.
   static MemDepResult getDef(Instruction *Inst) {
     assert(Inst && "Def requires inst");
     return MemDepResult(ValueTy::create<Def>(Inst));
   }
   static MemDepResult getClobber(Instruction *Inst) {
     assert(Inst && "Clobber requires inst");
     return MemDepResult(ValueTy::create<Clobber>(Inst));
   }
+  static MemDepResult getSelect(Instruction *Inst) {
+    return MemDepResult(ValueTy::create<Select>(Inst));
+  }
   static MemDepResult getNonLocal() {
     return MemDepResult(ValueTy::create<Other>(NonLocal));
   }
   static MemDepResult getNonFuncLocal() {
     return MemDepResult(ValueTy::create<Other>(NonFuncLocal));
   }
   static MemDepResult getUnknown() {
     return MemDepResult(ValueTy::create<Other>(Unknown));
   }
 
   /// Tests if this MemDepResult represents a query that is an instruction
   /// clobber dependency.
   bool isClobber() const { return Value.is<Clobber>(); }
 
   /// Tests if this MemDepResult represents a query that is an instruction
   /// definition dependency.
   bool isDef() const { return Value.is<Def>(); }
 
+  /// Tests if this MemDepResult represents a query that is an instruction
+  /// select dependency.
+  bool isSelect() const { return Value.is<Select>(); }
+
+  /// Tests if this MemDepResult represents a local query (Clobber/Def/Select).
+  bool isLocal() const { return isClobber() || isDef() || isSelect(); }
+
   /// Tests if this MemDepResult represents a query that is transparent to the
   /// start of the block, but where a non-local hasn't been done.
   bool isNonLocal() const {
     return Value.is<Other>() && Value.cast<Other>() == NonLocal;
   }
 
   /// Tests if this MemDepResult represents a query that is transparent to the
   /// start of the function.
public:
   Instruction *getInst() const {
     switch (Value.getTag()) {
     case Invalid:
       return Value.cast<Invalid>();
     case Clobber:
       return Value.cast<Clobber>();
     case Def:
       return Value.cast<Def>();
+    case Select:
+      return Value.cast<Select>();
     case Other:
       return nullptr;
     }
     llvm_unreachable("Unknown discriminant!");
   }
 
   bool operator==(const MemDepResult &M) const { return Value == M.Value; }
   bool operator!=(const MemDepResult &M) const { return Value != M.Value; }

llvm/lib/Analysis/MemoryDependenceAnalysis.cpp

while (ScanIt != BB->begin()) {
     // looking for a clobber in many cases; that's an alias property and is
     // handled by BasicAA.
     if (isa<AllocaInst>(Inst) || isNoAliasCall(Inst)) {
       const Value *AccessPtr = getUnderlyingObject(MemLoc.Ptr);
       if (AccessPtr == Inst || BatchAA.isMustAlias(Inst, AccessPtr))
         return MemDepResult::getDef(Inst);
     }
 
+    // If we found a select instruction for MemLoc pointer, return it as
+    // Select dependency.
+    if (isa<SelectInst>(Inst) && MemLoc.Ptr == Inst)
+      return MemDepResult::getSelect(Inst);
+
     if (isInvariantLoad)
       continue;
 
     // A release fence requires that all stores complete before it, but does
     // not prevent the reordering of following loads or stores 'before' the
     // fence.  As a result, we look past it when finding a dependency for
     // loads.  DSE uses this to find preceding stores to delete and thus we
     // can't bypass the fence if the query instruction is a store.
MemDepResult MemoryDependenceResults::getNonLocalInfoForBlock(
   if (ExistingResult)
     ExistingResult->setResult(Dep);
   else
     Cache->push_back(NonLocalDepEntry(BB, Dep));
 
   // If the block has a dependency (i.e. it isn't completely transparent to
   // the value), remember the reverse association because we just added it
   // to Cache!
-  if (!Dep.isDef() && !Dep.isClobber())
+  if (!Dep.isLocal())
     return Dep;
 
   // Keep the ReverseNonLocalPtrDeps map up to date so we can efficiently
   // update MemDep when we remove instructions.
   Instruction *Inst = Dep.getInst();
   assert(Inst && "Didn't depend on anything?");
   ValueIsLoadPair CacheKey(Loc.Ptr, isLoad);
   ReverseNonLocalPtrDeps[Inst].insert(CacheKey);

llvm/lib/Transforms/Scalar/GVN.cpp

 }
 
 /// Check if a load from pointer-select \p Address in \p DepBB can be converted
 /// to a value select. The following conditions need to be satisfied:
 /// 1. The pointer select (\p Address) must be defined in \p DepBB.
 /// 2. Both value operands of the pointer select must be loaded in the same
 /// basic block, before the pointer select.
 /// 3. There must be no instructions between the found loads and \p End that may
 /// clobber the loads.
 static std::optional<AvailableValue>
-tryToConvertLoadOfPtrSelect(BasicBlock *DepBB, BasicBlock::iterator End,
-                            Value *Address, Type *LoadTy, DominatorTree &DT,
+tryToConvertLoadOfPtrSelect(SelectInst *Sel, Type *LoadTy, DominatorTree &DT,
                             AAResults *AA) {
-
-  auto *Sel = dyn_cast_or_null<SelectInst>(Address);
-  if (!Sel || DepBB != Sel->getParent())
-    return std::nullopt;
-
   LoadInst *L1 = findDominatingLoad(Sel->getOperand(1), LoadTy, Sel, DT);
   LoadInst *L2 = findDominatingLoad(Sel->getOperand(2), LoadTy, Sel, DT);
   if (!L1 || !L2)
     return std::nullopt;
 
   // Ensure there are no accesses that may modify the locations referenced by
   // either L1 or L2 between L1, L2 and the specified End iterator.
   Instruction *EarlierLoad = L1->comesBefore(L2) ? L1 : L2;
   MemoryLocation L1Loc = MemoryLocation::get(L1);
   MemoryLocation L2Loc = MemoryLocation::get(L2);
-  if (any_of(make_range(EarlierLoad->getIterator(), End), [&](Instruction &I) {
+  if (any_of(make_range(EarlierLoad->getIterator(), Sel->getIterator()),
+             [&](Instruction &I) {
                return isModSet(AA->getModRefInfo(&I, L1Loc)) ||
                       isModSet(AA->getModRefInfo(&I, L2Loc));
              }))
     return std::nullopt;
 
   return AvailableValue::getSelect(Sel);
 }
 
 std::optional<AvailableValue>
 GVNPass::AnalyzeLoadAvailability(LoadInst *Load, MemDepResult DepInfo,
                                  Value *Address) {
-  if (!DepInfo.isDef() && !DepInfo.isClobber()) {
-    assert(isa<SelectInst>(Address));
-    return tryToConvertLoadOfPtrSelect(Load->getParent(), Load->getIterator(),
-                                       Address, Load->getType(),
-                                       getDominatorTree(), getAliasAnalysis());
-  }
-
-  assert((DepInfo.isDef() || DepInfo.isClobber()) &&
-         "expected a local dependence");
   assert(Load->isUnordered() && "rules below are incorrect for ordered access");
-
-  const DataLayout &DL = Load->getModule()->getDataLayout();
+  assert(DepInfo.isLocal() && "expected a local dependence");
 
   Instruction *DepInst = DepInfo.getInst();
+  if (DepInfo.isSelect()) {
+    // Check if load with Addr dependent from select can be converted to select
+    // between load values. There must be no instructions between the found
+    // loads and DepInst that may clobber the loads.
+    assert(isa<SelectInst>(DepInst) && "Dependency must be select instruction");
+    auto *Sel = cast<SelectInst>(DepInst);
+    return tryToConvertLoadOfPtrSelect(Sel, Load->getType(), getDominatorTree(),
+                                       getAliasAnalysis());
+  }
+
+  const DataLayout &DL = Load->getModule()->getDataLayout();
   if (DepInfo.isClobber()) {
     // If the dependence is to a store that writes to a superset of the bits
     // read by the load, we can extract the bits we need for the load from the
     // stored value.
     if (StoreInst *DepSI = dyn_cast<StoreInst>(DepInst)) {
       // Can't forward from non-atomic to atomic without violating memory model.
       if (Address && Load->isAtomic() <= DepSI->isAtomic()) {
         int Offset =
for (const auto &Dep : Deps) {
 
     if (DeadBlocks.count(DepBB)) {
       // Dead dependent mem-op disguise as a load evaluating the same value
       // as the load in question.
       ValuesPerBlock.push_back(AvailableValueInBlock::getUndef(DepBB));
       continue;
     }
 
-    // The address being loaded in this non-local block may not be the same as
-    // the pointer operand of the load if PHI translation occurs.  Make sure
-    // to consider the right address.
-    Value *Address = Dep.getAddress();
-
-    if (!DepInfo.isDef() && !DepInfo.isClobber()) {
-      if (auto R = tryToConvertLoadOfPtrSelect(
-              DepBB, DepBB->end(), Address, Load->getType(), getDominatorTree(),
-              getAliasAnalysis())) {
-        ValuesPerBlock.push_back(
-            AvailableValueInBlock::get(DepBB, std::move(*R)));
-        continue;
-      }
+    if (!DepInfo.isLocal()) {
       UnavailableBlocks.push_back(DepBB);
       continue;
     }
 
-    if (auto AV = AnalyzeLoadAvailability(Load, DepInfo, Address)) {
+    // The address being loaded in this non-local block may not be the same as
+    // the pointer operand of the load if PHI translation occurs.  Make sure
+    // to consider the right address.
+    if (auto AV = AnalyzeLoadAvailability(Load, DepInfo, Dep.getAddress())) {
       // subtlety: because we know this was a non-local dependency, we know
       // it's safe to materialize anywhere between the instruction within
       // DepInfo and the end of it's block.
       ValuesPerBlock.push_back(
           AvailableValueInBlock::get(DepBB, std::move(*AV)));
     } else {
       UnavailableBlocks.push_back(DepBB);
     }
bool GVNPass::processLoad(LoadInst *L) {
 
   // ... to a pointer that has been loaded from before...
   MemDepResult Dep = MD->getDependency(L);
 
   // If it is defined in another block, try harder.
   if (Dep.isNonLocal())
     return processNonLocalLoad(L);
 
-  Value *Address = L->getPointerOperand();
   // Only handle the local case below
-  if (!Dep.isDef() && !Dep.isClobber() && !isa<SelectInst>(Address)) {
+  if (!Dep.isLocal()) {
     // This might be a NonFuncLocal or an Unknown
     LLVM_DEBUG(
         // fast print dep, using operator<< on instruction is too slow.
         dbgs() << "GVN: load "; L->printAsOperand(dbgs());
         dbgs() << " has unknown dependence\n";);
     return false;
   }
 
-  auto AV = AnalyzeLoadAvailability(L, Dep, Address);
+  auto AV = AnalyzeLoadAvailability(L, Dep, L->getPointerOperand());
   if (!AV)
     return false;
 
   Value *AvailableValue = AV->MaterializeAdjustedValue(L, L, *this);
 
   // Replace the load!
   patchAndReplaceAllUsesWith(L, AvailableValue);
   markInstructionForDeletion(L);

llvm/test/Transforms/GVN/PRE/pre-load-through-select.ll

 ; CHECK-LABEL: @test_pointer_phi_select_simp_store_clobber_3(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br i1 [[COND:%.*]], label [[THEN:%.*]], label [[ELSE:%.*]]
 ; CHECK:       then:
 ; CHECK-NEXT:    store i32 99, ptr [[C:%.*]], align 4
 ; CHECK-NEXT:    [[L_1:%.*]] = load i32, ptr [[A:%.*]], align 4
 ; CHECK-NEXT:    [[L_2:%.*]] = load i32, ptr [[B:%.*]], align 4
 ; CHECK-NEXT:    [[CMP_I_I_I:%.*]] = icmp ult i32 [[L_1]], [[L_2]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[CMP_I_I_I]], i32 [[L_1]], i32 [[L_2]]
 ; CHECK-NEXT:    [[MIN_SELECT:%.*]] = select i1 [[CMP_I_I_I]], ptr [[A]], ptr [[B]]
 ; CHECK-NEXT:    br label [[EXIT:%.*]]
 ; CHECK:       else:
+; CHECK-NEXT:    [[RES_2_PRE:%.*]] = load i32, ptr [[A]], align 4
 ; CHECK-NEXT:    br label [[EXIT]]
 ; CHECK:       exit:
+; CHECK-NEXT:    [[RES_2:%.*]] = phi i32 [ [[SELECT]], [[THEN]] ], [ [[RES_2_PRE]], [[ELSE]] ]
 ; CHECK-NEXT:    [[P:%.*]] = phi ptr [ [[MIN_SELECT]], [[THEN]] ], [ [[A]], [[ELSE]] ]
-; CHECK-NEXT:    [[RES_2:%.*]] = load i32, ptr [[P]], align 4
 ; CHECK-NEXT:    ret i32 [[RES_2]]
 ;
 entry:
   br i1 %cond, label %then, label %else
 
 then:
   store i32 99, ptr %c
   %l.1 = load i32, ptr %a, align 4