[DSE] Make sure that DSE+MSSA can handle masked stores

[kparzysz]

Summary

Make changes to the MSSA-based parts of DSE to recognize and handle masked stores.

[kparzysz]

I think that the isMaskedStoreOverwrite could be merged into isOverwrite, but it'd need to take instructions in addition to the locations. I think it would be a bit cleaner than what this patch does.

[kparzysz]

Calling isMaskedStoreOverwrite from isOverwrite to avoid multiple checks.

Fixed a bug in ordering of NI and SI in the previous diff.

[fhahn]

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

414

Might be good to state that Later must overwrite Earlier. We could check that the mask of Later is a superset of the mask of Earlier. Might be good to add a TODO.

426

Wouldn't it be better to check if the pointers are must-alias?

1730–1733

Could you update the doc-comment?

1845

Is it possible to mark masked_store as IntrWriteMem? Then mayReadFromMemory should automatically return false for it.

[kparzysz]

Marked masked-storing intrinsics as write-only in 8b7c8f2c549d301fcea75d8e6e98a8ee160d5ff4.

[kparzysz]

Addressed comments.

[fhahn]

LGTM, thanks!

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

1851

unrelated change?

[kparzysz]

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

1851

Leftover from the "mayReadFromMemory" function from the previous diff. I'll remove it before committing.

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

enum OverwriteResult {
   OW_End,
   OW_PartialEarlierWithFullLater,
   OW_MaybePartial,
   OW_Unknown
 };
 
 } // end anonymous namespace
 
-/// Return 'OW_Complete' if a store to the 'Later' location completely
-/// overwrites a store to the 'Earlier' location. Return OW_MaybePartial
-/// if \p Later does not completely overwrite \p Earlier, but they both
-/// write to the same underlying object. In that case, use isPartialOverwrite to
-/// check if \p Later partially overwrites \p Earlier. Returns 'OW_Unknown' if
-/// nothing can be determined.
+/// Check if two instruction are masked stores that completely
+/// overwrite one another. More specifically, \p Later has to
+/// overwrite \p Earlier.
+template <typename AATy>
+static OverwriteResult isMaskedStoreOverwrite(const Instruction *Later,
+                                              const Instruction *Earlier,
+                                              AATy &AA) {
+  const auto *IIL = dyn_cast<IntrinsicInst>(Later);
+  const auto *IIE = dyn_cast<IntrinsicInst>(Earlier);
+  if (IIL == nullptr || IIE == nullptr)
+    return OW_Unknown;
+  if (IIL->getIntrinsicID() != Intrinsic::masked_store ||
+      IIE->getIntrinsicID() != Intrinsic::masked_store)
+    return OW_Unknown;
+  // Pointers.
+  Value *LP = IIL->getArgOperand(1)->stripPointerCasts();
+  Value *EP = IIE->getArgOperand(1)->stripPointerCasts();
+  if (LP != EP && !AA.isMustAlias(LP, EP))
+    return OW_Unknown;
+  // Masks.
+  // TODO: check that Later's mask is a superset of the Earlier's mask.
+  if (IIL->getArgOperand(3) != IIE->getArgOperand(3))
+    return OW_Unknown;
+  return OW_Complete;
+}
+
+/// Return 'OW_Complete' if a store to the 'Later' location (by \p LaterI
+/// instruction) completely overwrites a store to the 'Earlier' location.
+/// (by \p EarlierI instruction).
+/// Return OW_MaybePartial if \p Later does not completely overwrite
+/// \p Earlier, but they both write to the same underlying object. In that
+/// case, use isPartialOverwrite to check if \p Later partially overwrites
+/// \p Earlier. Returns 'OW_Unknown' if nothing can be determined.
 template <typename AATy>
 static OverwriteResult
-isOverwrite(const MemoryLocation &Later, const MemoryLocation &Earlier,
+isOverwrite(const Instruction *LaterI, const Instruction *EarlierI,
+            const MemoryLocation &Later, const MemoryLocation &Earlier,
             const DataLayout &DL, const TargetLibraryInfo &TLI,
             int64_t &EarlierOff, int64_t &LaterOff, AATy &AA,
             const Function *F) {
   // FIXME: Vet that this works for size upper-bounds. Seems unlikely that we'll
   // get imprecise values here, though (except for unknown sizes).
-  if (!Later.Size.isPrecise() || !Earlier.Size.isPrecise())
-    return OW_Unknown;
+  if (!Later.Size.isPrecise() || !Earlier.Size.isPrecise()) {
+    // Masked stores have imprecise locations, but we can reason about them
+    // to some extent.
+    return isMaskedStoreOverwrite(LaterI, EarlierI, AA);
+  }
 
   const uint64_t LaterSize = Later.Size.getValue();
   const uint64_t EarlierSize = Earlier.Size.getValue();
 
   const Value *P1 = Earlier.Ptr->stripPointerCasts();
   const Value *P2 = Later.Ptr->stripPointerCasts();
 
   // If the start pointers are the same, we just have to compare sizes to see if
if (EarlierOff >= LaterOff &&
       LaterSize >= EarlierSize &&
       uint64_t(EarlierOff - LaterOff) + EarlierSize <= LaterSize)
     return OW_Complete;
 
   // Later may overwrite earlier completely with other partial writes.
   return OW_MaybePartial;
 }
 
-static OverwriteResult isMaskedStoreOverwrite(Instruction *Later,
-                                              Instruction *Earlier) {
-  auto *IIL = dyn_cast<IntrinsicInst>(Later);
-  auto *IIE = dyn_cast<IntrinsicInst>(Earlier);
-  if (IIL == nullptr || IIE == nullptr)
-    return OW_Unknown;
-  if (IIL->getIntrinsicID() != Intrinsic::masked_store ||
-      IIE->getIntrinsicID() != Intrinsic::masked_store)
-    return OW_Unknown;
-  // Pointers.
-  if (IIL->getArgOperand(1) != IIE->getArgOperand(1))
-    return OW_Unknown;
-  // Masks.
-  if (IIL->getArgOperand(3) != IIE->getArgOperand(3))
-    return OW_Unknown;
-  return OW_Complete;
-}
-
 /// Return 'OW_Complete' if a store to the 'Later' location completely
 /// overwrites a store to the 'Earlier' location, 'OW_End' if the end of the
 /// 'Earlier' location is completely overwritten by 'Later', 'OW_Begin' if the
 /// beginning of the 'Earlier' location is overwritten by 'Later'.
 /// 'OW_PartialEarlierWithFullLater' means that an earlier (big) store was
 /// overwritten by a latter (smaller) store which doesn't write outside the big
 /// store's memory locations. Returns 'OW_Unknown' if nothing can be determined.
 /// NOTE: This function must only be called if both \p Later and \p Earlier
while (InstDep.isDef() || InstDep.isClobber()) {
       // If we find a write that is a) removable (i.e., non-volatile), b) is
       // completely obliterated by the store to 'Loc', and c) which we know that
       // 'Inst' doesn't load from, then we can remove it.
       // Also try to merge two stores if a later one only touches memory written
       // to by the earlier one.
       if (isRemovable(DepWrite) &&
           !isPossibleSelfRead(Inst, Loc, DepWrite, *TLI, *AA)) {
         int64_t InstWriteOffset, DepWriteOffset;
-        OverwriteResult OR = isOverwrite(Loc, DepLoc, DL, *TLI, DepWriteOffset,
-                                         InstWriteOffset, *AA, BB.getParent());
-        if (OR == OW_Unknown) {
-          // isOverwrite punts on MemoryLocations with an imprecise size, such
-          // as masked stores. Handle this here, somwewhat inelegantly.
-          OR = isMaskedStoreOverwrite(Inst, DepWrite);
-        }
+        OverwriteResult OR = isOverwrite(Inst, DepWrite, Loc, DepLoc, DL, *TLI,
+                                         DepWriteOffset, InstWriteOffset, *AA,
+                                         BB.getParent());
         if (OR == OW_MaybePartial)
           OR = isPartialOverwrite(Loc, DepLoc, DepWriteOffset, InstWriteOffset,
                                   DepWrite, IOL);
 
         if (OR == OW_Complete) {
           LLVM_DEBUG(dbgs() << "DSE: Remove Dead Store:\n  DEAD: " << *DepWrite
                             << "\n  KILLER: " << *Inst << '\n');
 
if (auto *CB = dyn_cast<CallBase>(I)) {
           return {MemoryLocation(CB->getArgOperand(0))};
         default:
           break;
         }
       }
       switch (CB->getIntrinsicID()) {
       case Intrinsic::init_trampoline:
         return {MemoryLocation(CB->getArgOperand(0))};
+      case Intrinsic::masked_store:
+        return {MemoryLocation::getForArgument(CB, 1, TLI)};
       default:
         break;
       }
       return None;
     }
 
     return MemoryLocation::getOrNone(I);
   }
 
-  /// Returns true if \p Use completely overwrites \p DefLoc.
-  bool isCompleteOverwrite(MemoryLocation DefLoc, Instruction *UseInst) {
+  /// Returns true if \p UseInst completely overwrites \p DefLoc
+  /// (stored by \p DefInst).
+  bool isCompleteOverwrite(MemoryLocation DefLoc, Instruction *DefInst,
+                           Instruction *UseInst) {
     // UseInst has a MemoryDef associated in MemorySSA. It's possible for a
     // MemoryDef to not write to memory, e.g. a volatile load is modeled as a
     // MemoryDef.
     if (!UseInst->mayWriteToMemory())
       return false;
 
     if (auto *CB = dyn_cast<CallBase>(UseInst))
       if (CB->onlyAccessesInaccessibleMemory())
         return false;
 
     int64_t InstWriteOffset, DepWriteOffset;
-    auto CC = getLocForWriteEx(UseInst);
-    return CC && isOverwrite(*CC, DefLoc, DL, TLI, DepWriteOffset,
-                             InstWriteOffset, BatchAA, &F) == OW_Complete;
+    if (auto CC = getLocForWriteEx(UseInst))
+      return isOverwrite(UseInst, DefInst, *CC, DefLoc, DL, TLI, DepWriteOffset,
+    return false;
   }
 
   /// Returns true if \p Def is not read before returning from the function.
   bool isWriteAtEndOfFunction(MemoryDef *Def) {
     LLVM_DEBUG(dbgs() << "  Check if def " << *Def << " ("
                       << *Def->getMemoryInst()
                       << ") is at the end the function \n");
 
if (MaybeTermLoc->second)
       DefLoc = MemoryLocation(getUnderlyingObject(DefLoc.Ptr));
     return BatchAA.isMustAlias(MaybeTermLoc->first, DefLoc);
   }
 
   // Returns true if \p Use may read from \p DefLoc.
   bool isReadClobber(MemoryLocation DefLoc, Instruction *UseInst) {
     // Monotonic or weaker atomic stores can be re-ordered and do not need to be
     // treated as read clobber.
     if (auto SI = dyn_cast<StoreInst>(UseInst))
       return isStrongerThan(SI->getOrdering(), AtomicOrdering::Monotonic);
 
     if (!UseInst->mayReadFromMemory())
       return false;
 
     if (auto *CB = dyn_cast<CallBase>(UseInst))
       if (CB->onlyAccessesInaccessibleMemory())
         return false;
 
     // NOTE: For calls, the number of stores removed could be slightly improved
     // by using AA.callCapturesBefore(UseInst, DefLoc, &DT), but that showed to
     // be expensive compared to the benefits in practice. For now, avoid more
     // expensive analysis to limit compile-time.
     return isRefSet(BatchAA.getModRefInfo(UseInst, DefLoc));
do {
         const Value *NIUnd = getUnderlyingObject(CurrentLoc->Ptr);
         if (DefUO != NIUnd) {
           StepAgain = true;
           Current = CurrentDef->getDefiningAccess();
         }
         continue;
       } else {
         int64_t InstWriteOffset, DepWriteOffset;
-        auto OR = isOverwrite(DefLoc, *CurrentLoc, DL, TLI, DepWriteOffset,
-                              InstWriteOffset, BatchAA, &F);
+        auto OR = isOverwrite(KillingI, CurrentI, DefLoc, *CurrentLoc, DL, TLI,
+                              DepWriteOffset, InstWriteOffset, BatchAA, &F);
         // If Current does not write to the same object as KillingDef, check
         // the next candidate.
         if (OR == OW_Unknown) {
           StepAgain = true;
           Current = CurrentDef->getDefiningAccess();
         } else if (OR == OW_MaybePartial) {
           // If KillingDef only partially overwrites Current, check the next
           // candidate if the partial step limit is exceeded. This aggressively
for (unsigned I = 0; I < WorkList.size(); I++) {
       // miss cases like the following
       //   1 = Def(LoE) ; <----- EarlierDef stores [0,1]
       //   2 = Def(1)   ; (2, 1) = NoAlias,   stores [2,3]
       //   Use(2)       ; MayAlias 2 *and* 1, loads [0, 3].
       //                  (The Use points to the *first* Def it may alias)
       //   3 = Def(1)   ; <---- Current  (3, 2) = NoAlias, (3,1) = MayAlias,
       //                  stores [0,1]
       if (MemoryDef *UseDef = dyn_cast<MemoryDef>(UseAccess)) {
-        if (isCompleteOverwrite(DefLoc, UseInst)) {
+        if (isCompleteOverwrite(DefLoc, KillingI, UseInst)) {
           if (!isInvisibleToCallerAfterRet(DefUO) &&
               UseAccess != EarlierAccess) {
             BasicBlock *MaybeKillingBlock = UseInst->getParent();
             if (PostOrderNumbers.find(MaybeKillingBlock)->second <
                 PostOrderNumbers.find(EarlierAccess->getBlock())->second) {
 
               LLVM_DEBUG(dbgs()
                          << "    ... found killing def " << *UseInst << "\n");
for (unsigned I = 0; I < ToCheck.size(); I++) {
                           << "\n  KILLER: " << *SI << '\n');
         State.deleteDeadInstruction(NI);
         ++NumFastStores;
         MadeChange = true;
       } else {
         // Check if NI overwrites SI.
         int64_t InstWriteOffset, DepWriteOffset;
         OverwriteResult OR =
-            isOverwrite(SILoc, NILoc, State.DL, TLI, DepWriteOffset,
+            isOverwrite(SI, NI, SILoc, NILoc, State.DL, TLI, DepWriteOffset,
                         InstWriteOffset, State.BatchAA, &F);
         if (OR == OW_MaybePartial) {
           auto Iter = State.IOLs.insert(
               std::make_pair<BasicBlock *, InstOverlapIntervalsTy>(
                   NI->getParent(), InstOverlapIntervalsTy()));
           auto &IOL = Iter.first->second;
           OR = isPartialOverwrite(SILoc, NILoc, DepWriteOffset, InstWriteOffset,
                                   NI, IOL);

llvm/test/Transforms/DeadStoreElimination/masked-dead-store.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -tbaa -dse -enable-dse-memoryssa=false -S < %s | FileCheck %s
+; RUN: opt -tbaa -dse -enable-dse-memoryssa=true -S < %s | FileCheck %s
 target datalayout = "e-m:e-p:32:32:32-a:0-n16:32-i64:64:64-i32:32:32-i16:16:16-i1:8:8-f32:32:32-f64:64:64-v32:32:32-v64:64:64-v512:512:512-v1024:1024:1024-v2048:2048:2048"
 
 define dllexport i32 @f0(i8** %a0, i8** %a1, i32 %a2, i32 %a3, i32 %a4, i32 %a5, i32 %a6, i32 %a7) #0 {
 ; CHECK-LABEL: @f0(
 ; CHECK-NEXT:  b0:
 ; CHECK-NEXT:    [[V0:%.*]] = getelementptr inbounds i8*, i8** [[A0:%.*]], i32 [[A2:%.*]]
 ; CHECK-NEXT:    [[V1:%.*]] = load i8*, i8** [[V0]], align 4, [[TBAA0:!tbaa !.*]]
 ; CHECK-NEXT:    [[V2:%.*]] = getelementptr i8, i8* [[V1]], i32 [[A3:%.*]]