[WIP][MemCpyOpt] implement liveness analysis based stack-move optimization

[khei4]

Summary

This is WIP
This patch extends multi-BB stack-move optimization using a liveness analysis framework as pcwalton does.
only focused on unescaped static allocas as the previous patch does,

This patch introduce BB/Instruction-level Liveness Analysis, basically with the same terminology with Modern Compiler Implementation book

1. capture check by CaptureTracker, collect whether all SSA uses are 1) LiveDef, which writes full size data to alloca (e.g. full-sized lifetime.start/end, memcpy/memset/memmove, store/load or 2) LiveUse, otherwise.
2. Compute the src/dest alloca Liveness, by LiveUse based backward approach.
   • If the BB's first SSA use is LiveUse of src/dest, this BB is LiveIn for src/dest.
   • The LiveIn BB's predecessors are LiveOut for the allocas.
     • The LiveIn BB's predecessors, which don't contain any LiveUse or LiveDef, are also LiveIn.

3. Check the BB-level liveness conflict; if it fails, check User-level liveness conflicts
   • BB-level
     • If the BB is also LiveIn or is LiveOut or contains LiveUse for src and dest simultaneously, then it conflict.
   • User-level
     • From the end of BB, propagate src and dest liveness (If met LiveUse, then live, otherwise LiveDef make it not live) (TODO: more proper explanation).

So basically above procedure is pcwalton's liveness analysis, but it miscompile and missed some opportunities as followings

misscompile (or miss opportunity, we can merge if we remove meaning less store %struct.Foo { i32 13, i32 13, i32 13 }, ptr %dest. These can be characterized by mergeable if we remove no Use, no LiveOut LiveDef `. But it requires more User-level checks if we try to merge them.

%src = alloca %struct.Foo, align 4
%dest = alloca %struct.Foo, align 4
call void @llvm.lifetime.start.p0(i64 12, ptr nocapture %src)
call void @llvm.lifetime.start.p0(i64 12, ptr nocapture %dest)
store %struct.Foo { i32 10, i32 20, i32 30 }, ptr %src

call void @llvm.memcpy.p0.p0.i64(ptr align 4 %dest, ptr align 4 %src, i64 12, i1 false)
store %struct.Foo { i32 13, i32 13, i32 13 }, ptr %dest
%1 = call i32 @use_nocapture(ptr nocapture %src)
call void @llvm.lifetime.end.p0(i64 12, ptr nocapture %src)
call void @llvm.lifetime.end.p0(i64 12, ptr nocapture %dest)
ret void

missed opportunity (if the Use is only read, no conflict happens.)

%src = alloca %struct.Foo, align 4
%dest = alloca %struct.Foo, align 4
call void @llvm.lifetime.start.p0(i64 12, ptr nocapture %src)
call void @llvm.lifetime.start.p0(i64 12, ptr nocapture %dest)
store %struct.Foo { i32 10, i32 20, i32 30 }, ptr %src

call void @llvm.memcpy.p0.p0.i64(ptr align 4 %dest, ptr align 4 %src, i64 12, i1 false)

%1 = call i32 @use_readonly(ptr nocapture %src)
%2 = call i32 @use_readonly(ptr nocapture %src)
%3 = call i32 @use_readonly(ptr nocapture %dest)
call void @llvm.lifetime.end.p0(i64 12, ptr nocapture %src)
call void @llvm.lifetime.end.p0(i64 12, ptr nocapture %dest)
ret void

Both cases can be covered by making LiveUse have ModRef information but balancing with compile time regressions.

TODO: Currently, the following will be miscompiled (src and dest would be merged)

%src = alloca %struct.Foo, align 4
%dest = alloca %struct.Foo, align 4
call void @llvm.lifetime.start.p0(i64 12, ptr nocapture %src)
call void @llvm.lifetime.start.p0(i64 12, ptr nocapture %dest)
store %struct.Foo { i32 10, i32 20, i32 30 }, ptr %src

call void @llvm.memcpy.p0.p0.i64(ptr align 4 %dest, ptr align 4 %src, i64 12, i1 false)

%1 = call i32 @use_nocapture(ptr nocapture %src)
%2 = call i32 @use_nocapture(ptr nocapture %dest)
call void @llvm.lifetime.end.p0(i64 12, ptr nocapture %src)
call void @llvm.lifetime.end.p0(i64 12, ptr nocapture %dest)
ret void

[nikic]

It looks like this is a diff against some intermediate version, rather than the full patch.

[khei4]

fix wrong diff

[khei4]

It looks like this is a diff against some intermediate version, rather than the full patch.

@nikic Thanks! I fixed and updated the description, although still WIP!

[khei4]

If I used the backward approach like this, I couldn't handle cleanly some cases that we could merge in the multi-BB case (for example, src and dest are read-only together; I believe it's a sufficiently familiar case for Rust). Still, I now noticed extending the ModRef analysis in a more forward flow-sensitive way could cover the use after def cases as complex as this patch is. So a change is planned.

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

bool MemCpyOptPass::performMemCpyToMemSetOptzn(MemCpyInst *MemCpy,
   return true;
 }
 
 // These helper classes are used for the stack-move optimization. See the
 // comments above performStackMoveOptzn() for more details.
 
 namespace {
 
-// TODO: rename
 enum Liveness { LiveUse, LiveDef };
 using LiveDefOrUseInst = PointerIntPair<Instruction *, 1, Liveness>;
 
 // Tracks liveness on the basic block level. This is conservative; see the
 // comments above performStackMoveOptzn() for justification.
 class BasicBlockLiveness {
 
   // Whether the alloca is live-in to the block (from predecessor basic blocks).
   bool LiveIn;
   // Whether the alloca is live-out from the block (to successor basic blocks).
   bool LiveOut;
   // Whether there's at least one use of the alloca in this basic block. This
   // flag is important for detecting liveness conflicts, since the other
   // information stored here isn't sufficient to determine that a use is present
   // if a definition precedes it.
   bool HasUse;
 
   // Type for the Instruction and LiveUse or LiveDef. if true, then LiveDef
   // The earliest definition or use we've seen, combined with the three bits
   // below.
   LiveDefOrUseInst FirstLiveUseOrDefInst;
 
   // All LiveUse or LiveDef Instructions on this block.
+  // FIXME: BSTMap for lookup and appearance order.
   SmallVector<LiveDefOrUseInst, 8> BBUsesOrDefs;
 
   // Records a first def or use instruction.
   void setFirstLiveDefOrUseInst(Instruction *I, Liveness L) {
     assert((!hasDefUseInst() ||
             I->comesBefore(getFirstLiveUseOrDefInst().getPointer())) &&
            "Tried to overwrite an earlier def or use with a later one!");
     FirstLiveUseOrDefInst.setPointer(I);
public:
   // Returns the earliest definition or use we've seen in this block.
   const LiveDefOrUseInst &getFirstLiveUseOrDefInst() const {
     return FirstLiveUseOrDefInst;
   }
 
   // Returns the all LiveUses and LiveDefs.
   SmallVector<LiveDefOrUseInst, 8> &getBBUsesOrDefs() { return BBUsesOrDefs; }
 
-  // Sort the all LiveUses and LiveDefs by appearance.
-  void sortBBUsesOrDefs() {
-    auto isLess = [](const LiveDefOrUseInst &L, const LiveDefOrUseInst &R) {
-      return L.getPointer()->comesBefore(R.getPointer());
-    };
-    std::sort(BBUsesOrDefs.begin(), BBUsesOrDefs.end(), isLess);
-  }
-
   // Returns true if there's a definition or use of the memory in this block.
   bool hasDefUseInst() const {
     return FirstLiveUseOrDefInst.getPointer() != nullptr;
   }
   // Returns true if the memory is live-in to this block (i.e. live-out of a
   // predecessor).
   bool isLiveIn() const { return LiveIn; }
   // Returns true if the memory is live-out of this block (i.e. live-in to a
public:
   void setLiveIn(bool On) { LiveIn = On; }
 
   // Adjusts the live-out flag for this block.
   void setLiveOut(bool On) { LiveOut = On; }
 
   // Records a new definition or use of the alloca being tracked within this
   // basic block.
   void update(Instruction *I, Liveness L) {
+    // Set use or def for the first instruction, BB is LiveIn if the first user
+    // is use.
     if (!hasDefUseInst() ||
         I->comesBefore(getFirstLiveUseOrDefInst().getPointer())) {
-      dbgs() << "set first use or def\n";
       setFirstLiveDefOrUseInst(I, L);
       setLiveIn(L == LiveUse);
     }
     if (L == LiveUse)
       setHasUse(true);
 
-    dbgs() << "add use or def\n";
     addLiveDefOrUseInst(I, L);
   }
 };
 
 } // namespace
 using BasicBlockLivenessMap = DenseMap<BasicBlock *, BasicBlockLiveness>;
 
 // Performs liveness dataflow analysis for an alloca at the basic block level as
static bool computeLiveness(BasicBlockLivenessMap &BBLivenessMap) {
 
   return true;
 }
 
 bool checkInstructionLevelLivenessConflict(
     BasicBlockLivenessMap &DestLivenessMap,
     BasicBlockLivenessMap &SrcLivenessMap, BasicBlock *BB, Instruction *Load,
     Instruction *Store) {
-  dbgs() << "hi from instruction level conflict on\n"
-         << BB->getNameOrAsOperand() << "\n";
-  // Check liveness inside the BB
+  // Check liveness inside the BB, check the uses in the order of appearance.
   auto SrcBBLiveness = SrcLivenessMap.lookup(BB);
   auto DestBBLiveness = DestLivenessMap.lookup(BB);
   bool SrcLive = SrcBBLiveness.isLiveOut();
   bool DestLive = DestBBLiveness.isLiveOut();
-  SrcBBLiveness.sortBBUsesOrDefs();
-  DestBBLiveness.sortBBUsesOrDefs();
   auto &SrcBBUsesOrDefs = SrcBBLiveness.getBBUsesOrDefs();
   auto &DestBBUsesOrDefs = DestBBLiveness.getBBUsesOrDefs();
-  LiveDefOrUseInst IL(nullptr);
-  bool IsSrc;
-  // This is destructive, but it's ok because there is no use for later.
-  while (!SrcBBUsesOrDefs.empty() || !DestBBUsesOrDefs.empty()) {
-    if (SrcBBUsesOrDefs.empty() ||
-        (!DestBBUsesOrDefs.empty() &&
-         SrcBBUsesOrDefs.back().getPointer()->comesBefore(
-             DestBBUsesOrDefs.back().getPointer()))) {
-      IL = DestBBUsesOrDefs.back();
-      DestBBUsesOrDefs.pop_back();
-      IsSrc = false;
-    } else {
-      IL = SrcBBUsesOrDefs.back();
-      SrcBBUsesOrDefs.pop_back();
-      IsSrc = true;
-    }
-
-    dbgs() << "use or def\n";
-    IL.getPointer()->print(dbgs());
-    dbgs() << "\n";
-    dbgs() << "DestLive=" << DestLive << "\n";
-    dbgs() << "SrcLive=" << SrcLive << "\n";
+  auto AllBBUsesOrDefs = SrcBBUsesOrDefs;
+  // FIXME: memcpy (Store, Load) is duplicated, so sorting is UB.
+  AllBBUsesOrDefs.insert(AllBBUsesOrDefs.end(), DestBBUsesOrDefs.begin(),
+                         DestBBUsesOrDefs.end());
+  auto isLess = [](const LiveDefOrUseInst &L, const LiveDefOrUseInst &R) {
+    return R.getPointer()->comesBefore(L.getPointer());
+  };
+  std::sort(AllBBUsesOrDefs.begin(), AllBBUsesOrDefs.end(), isLess);
+  for (auto IL : AllBBUsesOrDefs) {
     // FIXME: is this OK to exclude Load and store?
-    if (DestLive && SrcLive && IL.getPointer() != Load &&
-        IL.getPointer() != Store) {
+    if (DestLive && SrcLive) {
       LLVM_DEBUG(
           dbgs() << "Stack Move: Detected liveness conflict inside the basic "
                     "block \n  Instruction: "
                  << IL.getPointer() << "\n");
       return false;
     }
     Liveness L = IL.getInt();
-    if (IsSrc)
+    // FIXME: these equality check takes linear time, also failed because IL is
+    // not instruction but with LivenessInfo
+    if (std::find(SrcBBUsesOrDefs.begin(), SrcBBUsesOrDefs.end(), IL) !=
+        SrcBBUsesOrDefs.end())
       SrcLive = L == Liveness::LiveUse;
-    else
+
+    if (std::find(DestBBUsesOrDefs.begin(), DestBBUsesOrDefs.end(), IL) !=
+        DestBBUsesOrDefs.end())
       DestLive = L == Liveness::LiveUse;
   }
   return true;
 }
 
 bool checkBBLevelLivenessConflict(BasicBlockLivenessMap &DestLivenessMap,
                                   BasicBlockLivenessMap &SrcLivenessMap,
                                   Instruction *Load, Instruction *Store) {
-
-  dbgs() << "hi from bb level conflict\n";
-  dbgs() << "src size is\n";
-  dbgs() << SrcLivenessMap.size();
-  dbgs() << "\nDest size is\n";
-  dbgs() << DestLivenessMap.size();
   // Check for liveness conflicts on the basic block level (with the exception
   // of the basic block containing the memcpy).
   for (const auto &[BB, DestLiveness] : DestLivenessMap) {
-    dbgs() << "Basic Block: " << BB->getNameOrAsOperand() << "\n";
-    dbgs() << "islivenayware: " << DestLiveness.isLiveAnywhereOrHasUses()
-           << "\n";
-    dbgs() << "corresponding src islivenayware: "
-           << SrcLivenessMap.lookup(BB).isLiveAnywhereOrHasUses() << "\n";
-
     if (DestLiveness.isLiveAnywhereOrHasUses() &&
         SrcLivenessMap.lookup(BB).isLiveAnywhereOrHasUses()) {
       LLVM_DEBUG(
           dbgs() << "Stack Move: Detected Basic Block liveness conflict, "
                     "\n  Basic Block: "
                  << BB->getNameOrAsOperand() << "\n");
       if (!checkInstructionLevelLivenessConflict(
               DestLivenessMap, SrcLivenessMap, BB, Load, Store))
while (!Worklist.empty()) {
               << "Stack Move: Exceeded max uses to see ModRef, bailing\n");
           return false;
         }
         if (!Visited.insert(&U).second)
           continue;
         switch (DetermineUseCaptureKind(U, IsDereferenceableOrNull)) {
         case UseCaptureKind::MAY_CAPTURE:
           LLVM_DEBUG({
-            dbgs() << "Stack Move: Alloca was captured:";
+            dbgs() << "Stack Move: Alloca may be captured:";
             dbgs() << "\n" << U;
             dbgs() << "\n";
           });
           return false;
         case UseCaptureKind::PASSTHROUGH:
           // FIXME: should we check liveness on this? I guess not, because these
           // are generally ignorable like gep.
           // Instructions cannot have non-instruction users.
           Worklist.push_back(cast<Instruction>(U.getUser()));
           continue;
         case UseCaptureKind::NO_CAPTURE: {
           auto *UI = cast<Instruction>(U.getUser());
           if (!DomI) {
             PDom = DomI = UI;
           } else {
             DomI = DT->findNearestCommonDominator(DomI, UI);
             if (PDom)
               PDom = findNearestCommonPostDominator(PDom, UI, PDT);
           }
 
-          dbgs() << "this is nocapture\n";
-          UI->print(dbgs());
-          dbgs() << "\n";
           // If an instruction overwrites all bytes of the alloca, it's a
           // definition, not a use. Detect those cases here.
           if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(UI)) {
             if (II->isLifetimeStartOrEnd()) {
               // We treat a call to a lifetime intrinsic that covers the entire
               // alloca as a definition, since both llvm.lifetime.start and
               // llvm.lifetime.end intrinsics conceptually fill all the bytes of
               // the alloca with an undefined value. We also note these
               // locations of these intrinsic calls so that we can delete them
               // later if the optimization succeeds.
               int64_t Size =
                   cast<ConstantInt>(II->getArgOperand(0))->getSExtValue();
 
-              dbgs() << "hello from ll capture tracking\n";
               if (Size < 0 || uint64_t(Size) == SrcSize) {
-                dbgs() << "lifeitme is counted as LiveDef\n";
-                II->print(dbgs());
                 BBLivenessMap[II->getParent()].update(II, Liveness::LiveDef);
                 LifetimeMarkers.push_back(II);
                 continue;
               }
             } else if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(II)) {
               if (MI->getArgOperandNo(&U) == 0) {
                 if (ConstantInt *CI = dyn_cast<ConstantInt>(MI->getLength())) {
                   if (CI->getZExtValue() == (*SrcSize).getFixedValue()) {
 
-                    dbgs() << "hi from memintrinsic def\n";
-                    MI->print(dbgs());
-                    dbgs() << "\n";
                     // Memcpy, memmove, and memset instructions that fill every
                     // byte of the alloca are definitions.
                     BBLivenessMap[MI->getParent()].update(MI,
                                                           Liveness::LiveDef);
                     // TODO: remove dupulication of the noalias insertion
                     if (MI->hasMetadata(LLVMContext::MD_noalias))
                       NoAliasInstrs.insert(MI);
                     continue;
                   }
                 }
               }
             }
           } else if (StoreInst *SI = dyn_cast<StoreInst>(UI)) {
             // Stores that overwrite all bytes of the alloca are definitions.
             if (U.getOperandNo() == 1 &&
                 DL.getTypeStoreSize(SI->getValueOperand()->getType()) ==
                     (*SrcSize).getFixedValue()) {
               BBLivenessMap[SI->getParent()].update(SI, Liveness::LiveDef);
               if (SI->hasMetadata(LLVMContext::MD_noalias))
                 NoAliasInstrs.insert(SI);
               continue;
             }
           }
-          dbgs() << "BB found live use\n";
-          UI->print(dbgs());
-          dbgs() << "\n";
           BBLivenessMap[UI->getParent()].update(UI, Liveness::LiveUse);
           if (UI->hasMetadata(LLVMContext::MD_noalias))
             NoAliasInstrs.insert(UI);
         }
         }
       }
     }
     return true;