[clang] [llvm] [Transforms] Add LifetimeMovePass (PR #144319)

[Tomasz Miąsko via llvm-commits]

================
@@ -0,0 +1,341 @@
+//===- LifetimeMove.cpp - Narrowing lifetimes -----------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// The LifetimeMovePass identifies the precise lifetime range of allocas and
+// repositions lifetime markers to stricter positions.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/LifetimeMove.h"
+#include "llvm/Analysis/CFG.h"
+#include "llvm/Analysis/CaptureTracking.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/PostDominators.h"
+#include "llvm/Analysis/PtrUseVisitor.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/Transforms/Coroutines/CoroInstr.h"
+
+#define DEBUG_TYPE "lifetime-move"
+
+namespace llvm {
+namespace {
+class LifetimeMover : public PtrUseVisitor<LifetimeMover> {
+  using This = LifetimeMover;
+  using Base = PtrUseVisitor<LifetimeMover>;
+
+  const DominatorTree &DT;
+  const LoopInfo &LI;
+
+  SmallVector<AllocaInst *, 4> Allocas;
+  // Critical points are instructions where the crossing of a variable's
+  // lifetime makes a difference. We attempt to rise lifetime.end
+  // before critical points and sink lifetime.start after them.
+  SmallVector<Instruction *, 4> CriticalPoints;
+
+  SmallVector<Instruction *, 2> LifetimeStarts;
+  SmallVector<Instruction *, 2> LifetimeEnds;
+  SmallVector<Instruction *, 8> OtherUsers;
+  SmallPtrSet<BasicBlock *, 2> LifetimeStartBBs;
+  SmallPtrSet<BasicBlock *, 2> UserBBs;
+
+public:
+  LifetimeMover(Function &F, const DominatorTree &DT, const LoopInfo &LI);
+
+  bool run();
+
+  void visitInstruction(Instruction &I);
+  void visitPHINode(PHINode &I);
+  void visitSelectInst(SelectInst &I);
+  void visitStoreInst(StoreInst &SI);
+  void visitIntrinsicInst(IntrinsicInst &II);
+  void visitMemIntrinsic(MemIntrinsic &I);
+  void visitCallBase(CallBase &CB);
+
+private:
+  bool sinkLifetimeStartMarkers(AllocaInst *AI);
+  bool riseLifetimeEndMarkers();
+  void reset();
+};
+} // namespace
+
+LifetimeMover::LifetimeMover(Function &F, const DominatorTree &DT,
+                             const LoopInfo &LI)
+    : Base(F.getDataLayout()), DT(DT), LI(LI) {
+  for (Instruction &I : instructions(F)) {
+    if (auto *AI = dyn_cast<AllocaInst>(&I))
+      Allocas.push_back(AI);
+    else if (isa<LifetimeIntrinsic>(I))
+      continue;
+    else if (isa<AnyCoroSuspendInst>(I))
+      CriticalPoints.push_back(&I);
+    else if (isa<CallInst>(I))
+      CriticalPoints.push_back(&I);
+    else if (isa<InvokeInst>(I))
+      CriticalPoints.push_back(&I);
+  }
+}
+
+bool LifetimeMover::run() {
+  bool Changed = false;
+  for (auto *AI : Allocas) {
+    reset();
+    Base::visitPtr(*AI);
+
+    if (!LifetimeStarts.empty())
+      Changed |= sinkLifetimeStartMarkers(AI);
+
+    // Do not move lifetime.end if alloca escapes
+    if (!LifetimeEnds.empty() && !PI.isEscaped())
+      Changed |= riseLifetimeEndMarkers();
----------------
tmiasko wrote:

> > There are no executions where allocas are live at the same time and their storage overlaps.
> 
> You mean 'does not overlap'?

I meant what I wrote. There is a negation in the sentence already.

> The community has not yet reached a consensus on comparing a pointer outside its lifetime range, according to the RFC above. Therefore, the input IR should be considered invalid. What do we expect transform passes do on an invalid input? 

It is perfectly valid. LLVM Language Reference Manual explains:

> It is undefined behavior to access an allocated object that isn't alive, but operations that don't dereference it such as `getelementptr`, `ptrtoint` and `icmp` return a valid result. This explains code motion of these instructions across operations that impact the object's lifetime. 

This is mostly irrelevant to the example anyway, since it can be easily rewritten so that all references to alloca are within its lifetime.

> Specifically for the demo, LifetimeMovePass does nothing.

The example of course needs to be placed in a coroutine and extended with critical points, but LifetimeMovePass does pretty much identical transformation. For example in https://alive2.llvm.org/ce/z/CZQAw7, `tgt` is a result of running the transformation on `src` obtained from `opt -passes=lifetime-move`.

I think Alive2 quite correctly points out that this changes observable behavior. 

https://github.com/llvm/llvm-project/pull/144319