[CodeExtractor] Replace uses of extracted bitcasts in out-of-region lifetime markers

[ggeorgakoudis]

Summary

CodeExtractor handles bitcasts in the extracted region that have
lifetime markers users in the outer region as outputs. That
creates unnecessary alloca/reload instructions and extra lifetime
markers. The patch identifies those cases, and replaces uses in
out-of-region lifetime markers with new bitcasts in the outer region.

Example

define void @foo() {
entry:
  %0 = alloca i32
  br label %extract

extract:
  %1 = bitcast i32* %0 to i8*
  call void @llvm.lifetime.start.p0i8(i64 4, i8* %1)
  call void @use(i32* %0)
  br label %exit

exit:
  call void @use(i32* %0)
  call void @llvm.lifetime.end.p0i8(i64 4, i8* %1)
  ret void
}

Current extraction

define void @foo() {
entry:
  %.loc = alloca i8*, align 8
  %0 = alloca i32, align 4
  br label %codeRepl

codeRepl:                                         ; preds = %entry
  %lt.cast = bitcast i8** %.loc to i8*
  call void @llvm.lifetime.start.p0i8(i64 -1, i8* %lt.cast)
  %lt.cast1 = bitcast i32* %0 to i8*
  call void @llvm.lifetime.start.p0i8(i64 -1, i8* %lt.cast1)
  call void @foo.extract(i32* %0, i8** %.loc)
  %.reload = load i8*, i8** %.loc, align 8
  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %lt.cast)
  br label %exit

exit:                                             ; preds = %codeRepl
  call void @use(i32* %0)
  call void @llvm.lifetime.end.p0i8(i64 4, i8* %.reload)
  ret void
}

define internal void @foo.extract(i32* %0, i8** %.out) {
newFuncRoot:
  br label %extract

exit.exitStub:                                    ; preds = %extract
  ret void

extract:                                          ; preds = %newFuncRoot
  %1 = bitcast i32* %0 to i8*
  store i8* %1, i8** %.out, align 8
  call void @use(i32* %0)
  br label %exit.exitStub
}

Extraction with patch

define void @foo() {
entry:
  %0 = alloca i32, align 4
  br label %codeRepl

codeRepl:                                         ; preds = %entry
  %lt.cast1 = bitcast i32* %0 to i8*
  call void @llvm.lifetime.start.p0i8(i64 -1, i8* %lt.cast1)
  call void @foo.extract(i32* %0)
  br label %exit

exit:                                             ; preds = %codeRepl
  call void @use(i32* %0)
  %lt.cast = bitcast i32* %0 to i8*
  call void @llvm.lifetime.end.p0i8(i64 4, i8* %lt.cast)
  ret void
}

define internal void @foo.extract(i32* %0) {
newFuncRoot:
  br label %extract

exit.exitStub:                                    ; preds = %extract
  ret void

extract:                                          ; preds = %newFuncRoot
  %1 = bitcast i32* %0 to i8*
  call void @use(i32* %0)
  br label %exit.exitStub
}

[ggeorgakoudis]

Format better the patch message

[vsk]

In the example, what pass/component is responsible for creating %1 = bitcast i32* %0 to i8*?

If it's CodeExtractor itself, it would be nice to solve the problem earlier by getting rid of bitcasts which only exist to wire up lifetime markers. I tried doing that in 26ee8aff2, but had to revert (099bffe7f) because llvm doesn't support specializing its lifetime intrinsics on opaque pointer types. That's a bug, I think: if we fix that, perhaps this whole problem goes away.

OTOH, if this %1 = bitcast i32* %0 to i8* value comes from a frontend or another pass, I think the approach taken in this patch is reasonable. Please add a test (perhaps in the CodeExtractor unittest).

[ggeorgakoudis]

In D90689#2371522, @vsk wrote:

In the example, what pass/component is responsible for creating %1 = bitcast i32* %0 to i8*?

If it's CodeExtractor itself, it would be nice to solve the problem earlier by getting rid of bitcasts which only exist to wire up lifetime markers. I tried doing that in 26ee8aff2, but had to revert (099bffe7f) because llvm doesn't support specializing its lifetime intrinsics on opaque pointer types. That's a bug, I think: if we fix that, perhaps this whole problem goes away.

OTOH, if this %1 = bitcast i32* %0 to i8* value comes from a frontend or another pass, I think the approach taken in this patch is reasonable. Please add a test (perhaps in the CodeExtractor unittest).

Many thanks for the quick reply @vsk! The bitcast is generated by the frontend, not by CodeExtractor. Our use case is outlining a merged parallel region that includes sequential code which include bitcasts like that. I am going to add a unit test.

I have two questions about the CodeExtractor.

1. I have second thoughts on changing the IR within the findAllocas. Should I create a helper data structure, e.g., ReplaceBitcastUses as in for SinkCands, and do the transformation inside the extractCodeRegion caller?
2. Is the assumption that findAllocas is always called before findInputsOutputs? If findAllocas is not called, my concern is that findInputsOutputs will report an output on the use of the bitcast although this output would be removed if findAllocas has been called. This problem does not appear for extractCodeRegion because it always calls findAllocas but there may be other modules using the CodeExtractor interface. I could change findInputsOutputs not to include bitcast uses to lifetime markers as outputs regardless of whether findAllocas has been called. Is that reasonable or does that create other problems?

Thanks again!

[vsk]

In D90689#2371689, @ggeorgakoudis wrote:

In D90689#2371522, @vsk wrote:

In the example, what pass/component is responsible for creating %1 = bitcast i32* %0 to i8*?

If it's CodeExtractor itself, it would be nice to solve the problem earlier by getting rid of bitcasts which only exist to wire up lifetime markers. I tried doing that in 26ee8aff2, but had to revert (099bffe7f) because llvm doesn't support specializing its lifetime intrinsics on opaque pointer types. That's a bug, I think: if we fix that, perhaps this whole problem goes away.

OTOH, if this %1 = bitcast i32* %0 to i8* value comes from a frontend or another pass, I think the approach taken in this patch is reasonable. Please add a test (perhaps in the CodeExtractor unittest).

Many thanks for the quick reply @vsk! The bitcast is generated by the frontend, not by CodeExtractor. Our use case is outlining a merged parallel region that includes sequential code which include bitcasts like that. I am going to add a unit test.

I have two questions about the CodeExtractor.

1. I have second thoughts on changing the IR within the findAllocas. Should I create a helper data structure, e.g., ReplaceBitcastUses as in for SinkCands, and do the transformation inside the extractCodeRegion caller?
2. Is the assumption that findAllocas is always called before findInputsOutputs? If findAllocas is not called, my concern is that findInputsOutputs will report an output on the use of the bitcast although this output would be removed if findAllocas has been called. This problem does not appear for extractCodeRegion because it always calls findAllocas but there may be other modules using the CodeExtractor interface. I could change findInputsOutputs not to include bitcast uses to lifetime markers as outputs regardless of whether findAllocas has been called. Is that reasonable or does that create other problems?

The status quo is that findAllocas destructively updates the IR in ways that affect the results of findInputsOutputs: you're correct that this dependency can create issues. Imho CodeExtractor intermixes analysis and transformation too much. This makes it hard to improve its transformations (evidenced by your question (1)); it also makes it hard for client code to analyze extraction profitability, since you need to do the extraction to figure out what the real input/output set will be. (Before extractCodeRegions, you can query findInputsOutputs, but it's just guessing.)

Taking a short-term view, having findAllocas perform a new destructive update would get the job done and be in keeping with the current design. In the long term, I think it'd be nice to have a redesigned CodeExtractor that keeps analysis and transformation neatly separated. (that said, I'm not sure how best to achieve that (.. and I'm no longer actively working on it).)

[ggeorgakoudis]

Add unit test

[ggeorgakoudis]

I see your point @vsk. I have added a unit test. Anything else missing for acceptance?

[vsk]

Thanks, lgtm.

llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp

327

nit, it's a bit more idiomatic to write EXPECT_EQ(outputs.size(), 0U).

llvm/lib/Transforms/Utils/CodeExtractor.cpp

for (AllocaInst *AI : CEAC.getAllocas()) {
     LifetimeMarkerInfo MarkerInfo = getLifetimeMarkers(CEAC, AI, ExitBlock);
     bool Moved = moveOrIgnoreLifetimeMarkers(MarkerInfo);
     if (Moved) {
       LLVM_DEBUG(dbgs() << "Sinking alloca: " << *AI << "\n");
       SinkCands.insert(AI);
       continue;
     }
 
+    // Find bitcasts in the outlined region that have lifetime marker users
+    // outside that region. Replace the lifetime marker use with an
+    // outside region bitcast to avoid unnecessary alloca/reload instructions
+    // and extra lifetime markers.
+    SmallVector<Instruction *, 2> LifetimeBitcastUsers;
+    for (User *U : AI->users()) {
+      if (!definedInRegion(Blocks, U))
+        continue;
+
+      if (U->stripInBoundsConstantOffsets() != AI)
+        continue;
+
+      Instruction *Bitcast = cast<Instruction>(U);
+      for (User *BU : Bitcast->users()) {
+        IntrinsicInst *IntrInst = dyn_cast<IntrinsicInst>(BU);
+        if (!IntrInst)
+          continue;
+
+        if (!IntrInst->isLifetimeStartOrEnd())
+          continue;
+
+        if (definedInRegion(Blocks, IntrInst))
+          continue;
+
+        LLVM_DEBUG(dbgs() << "Replace use of extracted region bitcast"
+                          << *Bitcast << " in out-of-region lifetime marker "
+                          << *IntrInst << "\n");
+        LifetimeBitcastUsers.push_back(IntrInst);
+      }
+    }
+
+    for (Instruction *I : LifetimeBitcastUsers) {
+      Module *M = AIFunc->getParent();
+      LLVMContext &Ctx = M->getContext();
+      auto Int8PtrTy = Type::getInt8PtrTy(Ctx);
+      CastInst *CastI =
+          CastInst::CreatePointerCast(AI, Int8PtrTy, "lt.cast", I);
+      I->replaceUsesOfWith(I->getOperand(1), CastI);
+    }
+
     // Follow any bitcasts.
     SmallVector<Instruction *, 2> Bitcasts;
     SmallVector<LifetimeMarkerInfo, 2> BitcastLifetimeInfo;
     for (User *U : AI->users()) {
       if (U->stripInBoundsConstantOffsets() == AI) {
         Instruction *Bitcast = cast<Instruction>(U);
         LifetimeMarkerInfo LMI = getLifetimeMarkers(CEAC, Bitcast, ExitBlock);
         if (LMI.LifeStart) {

llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp

TEST(CodeExtractor, ExtractAndInvalidateAssumptionCache) {
 
   CodeExtractorAnalysisCache CEAC(*Func);
   Function *Outlined = CE.extractCodeRegion(CEAC);
   EXPECT_TRUE(Outlined);
   EXPECT_FALSE(verifyFunction(*Outlined));
   EXPECT_FALSE(verifyFunction(*Func));
   EXPECT_FALSE(CE.verifyAssumptionCache(*Func, *Outlined, &AC));
 }
+
+TEST(CodeExtractor, RemoveBitcastUsesFromOuterLifetimeMarkers) {
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M(parseAssemblyString(R"ir(
+    target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+    target triple = "x86_64-unknown-linux-gnu"
+
+    declare void @use(i32*)
+    declare void @llvm.lifetime.start.p0i8(i64, i8*)
+    declare void @llvm.lifetime.end.p0i8(i64, i8*)
+
+    define void @foo() {
+    entry:
+      %0 = alloca i32
+      br label %extract
+
+    extract:
+      %1 = bitcast i32* %0 to i8*
+      call void @llvm.lifetime.start.p0i8(i64 4, i8* %1)
+      call void @use(i32* %0)
+      br label %exit
+
+    exit:
+      call void @use(i32* %0)
+      call void @llvm.lifetime.end.p0i8(i64 4, i8* %1)
+      ret void
+    }
+  )ir",
+                                                Err, Ctx));
+
+  Function *Func = M->getFunction("foo");
+  SmallVector<BasicBlock *, 1> Blocks{getBlockByName(Func, "extract")};
+
+  CodeExtractor CE(Blocks);
+  EXPECT_TRUE(CE.isEligible());
+
+  CodeExtractorAnalysisCache CEAC(*Func);
+  SetVector<Value *> inputs, outputs, SinkingCands, HoistingCands;
+  BasicBlock *CommonExit = nullptr;
+  CE.findAllocas(CEAC, SinkingCands, HoistingCands, CommonExit);
+  CE.findInputsOutputs(inputs, outputs, SinkingCands);
+  EXPECT_TRUE(outputs.size() == 0);
+
+  Function *Outlined = CE.extractCodeRegion(CEAC);
+  EXPECT_TRUE(Outlined);
+  EXPECT_FALSE(verifyFunction(*Outlined));
+  EXPECT_FALSE(verifyFunction(*Func));
+}
 } // end anonymous namespace