[llvm] [Extractor] Use function return for the one and only output (PR #191824)

Shimin Cui via llvm-commits llvm-commits at lists.llvm.org
Mon Apr 13 07:30:15 PDT 2026


https://github.com/scui-ibm created https://github.com/llvm/llvm-project/pull/191824

Currently code extractor uses parameters to pass outputs.  Alloca/store/load instructions are used to get the output value in the parent functions. 

When there is only one output from the extracted code (this is one of the most common cases), using the function return for the only one output can facilitate the other transformations (eg, tail call opt). This is to modify the code for the extracted function to return the output value if there is only one output for the extracted region. 

>From 8bab39ea3ac9e6633a82c7655a13780cb5e360bf Mon Sep 17 00:00:00 2001
From: Shimin Cui <scui at ca.ibm.com>
Date: Mon, 13 Apr 2026 14:11:58 +0000
Subject: [PATCH] [Extractor] Use function return for the one and only output

---
 .../llvm/Transforms/Utils/CodeExtractor.h     |  18 ++-
 llvm/lib/Transforms/IPO/BlockExtractor.cpp    |   5 +-
 llvm/lib/Transforms/IPO/IROutliner.cpp        |   4 +-
 llvm/lib/Transforms/Utils/CodeExtractor.cpp   |  37 ++++--
 .../CodeExtractor/PartialInlineAnd.ll         |   2 +-
 .../CodeExtractor/PartialInlineAttributes.ll  |   8 +-
 .../CodeExtractor/PartialInlineDebug.ll       |   8 +-
 .../PartialInlineInvokeProducesOutVal.ll      |  11 +-
 .../CodeExtractor/PartialInlineOrAnd.ll       |   2 +-
 .../PartialInlineVarArgsDebug.ll              |   4 +-
 .../HotColdSplit/apply-penalty-for-outputs.ll |   6 +-
 .../HotColdSplit/duplicate-phi-preds-crash.ll |   4 +-
 .../phi-with-distinct-outlined-values.ll      |   5 +-
 .../HotColdSplit/split-phis-in-exit-blocks.ll |   8 +-
 .../HotColdSplit/succ-block-with-self-edge.ll |   8 +-
 .../Transforms/Utils/CodeExtractorTest.cpp    | 108 +++++++++++++++---
 16 files changed, 171 insertions(+), 67 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Utils/CodeExtractor.h b/llvm/include/llvm/Transforms/Utils/CodeExtractor.h
index e117a33d3435f..b86e0a00d9c80 100644
--- a/llvm/include/llvm/Transforms/Utils/CodeExtractor.h
+++ b/llvm/include/llvm/Transforms/Utils/CodeExtractor.h
@@ -132,6 +132,13 @@ class LLVM_ABI CodeExtractor {
   // space.
   bool ArgsInZeroAddressSpace;
 
+  // If true, the outlined function always return void even when there is only
+  // one output.
+  bool VoidReturnWithSingleOutput;
+
+  // If set, the return value of the outline function.
+  Value *FuncRetVal = nullptr;
+
 public:
   /// Create a code extractor for a sequence of blocks.
   ///
@@ -147,13 +154,16 @@ class LLVM_ABI CodeExtractor {
   /// which case it will be placed in the entry block of the function from which
   /// the code is being extracted. If ArgsInZeroAddressSpace param is set to
   /// true, then the aggregate param pointer of the outlined function is
-  /// declared in zero address space.
+  /// declared in zero address space. If VoidReturnWithSingleOutput is set to
+  /// true, then the return type of the outlined function is set void even if
+  /// there is only one output.
   CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
                 bool AggregateArgs = false, BlockFrequencyInfo *BFI = nullptr,
                 BranchProbabilityInfo *BPI = nullptr,
                 AssumptionCache *AC = nullptr, bool AllowVarArgs = false,
                 bool AllowAlloca = false, BasicBlock *AllocationBlock = nullptr,
-                std::string Suffix = "", bool ArgsInZeroAddressSpace = false);
+                std::string Suffix = "", bool ArgsInZeroAddressSpace = false,
+                bool VoidReturnWithSingleOutput = false);
 
   /// Perform the extraction, returning the new function.
   ///
@@ -201,7 +211,7 @@ class LLVM_ABI CodeExtractor {
   /// on the cost however.
   void findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
                          const ValueSet &Allocas,
-                         bool CollectGlobalInputs = false) const;
+                         bool CollectGlobalInputs = false);
 
   /// Check if life time marker nodes can be hoisted/sunk into the outline
   /// region.
@@ -300,7 +310,7 @@ class LLVM_ABI CodeExtractor {
   /// into the original function's control flow.
   void
   insertReplacerCall(Function *oldFunction, BasicBlock *header,
-                     BasicBlock *codeReplacer, const ValueSet &outputs,
+                     CallInst *ReplacerCall, const ValueSet &outputs,
                      ArrayRef<Value *> Reloads,
                      const DenseMap<BasicBlock *, BlockFrequency> &ExitWeights);
 };
diff --git a/llvm/lib/Transforms/IPO/BlockExtractor.cpp b/llvm/lib/Transforms/IPO/BlockExtractor.cpp
index 23991afc511a0..52a20227da633 100644
--- a/llvm/lib/Transforms/IPO/BlockExtractor.cpp
+++ b/llvm/lib/Transforms/IPO/BlockExtractor.cpp
@@ -168,7 +168,10 @@ bool BlockExtractor::runOnModule(Module &M) {
       Changed = true;
     }
     CodeExtractorAnalysisCache CEAC(*BBs[0]->getParent());
-    Function *F = CodeExtractor(BlocksToExtractVec).extractCodeRegion(CEAC);
+    Function *F =
+        CodeExtractor(BlocksToExtractVec, nullptr, false, nullptr, nullptr,
+                      nullptr, false, false, nullptr, "", false, true)
+            .extractCodeRegion(CEAC);
     if (F)
       LLVM_DEBUG(dbgs() << "Extracted group '" << (*BBs.begin())->getName()
                         << "' in: " << F->getName() << '\n');
diff --git a/llvm/lib/Transforms/IPO/IROutliner.cpp b/llvm/lib/Transforms/IPO/IROutliner.cpp
index 41c036496487c..08a77993423af 100644
--- a/llvm/lib/Transforms/IPO/IROutliner.cpp
+++ b/llvm/lib/Transforms/IPO/IROutliner.cpp
@@ -2789,7 +2789,7 @@ unsigned IROutliner::doOutline(Module &M) {
       OS->Candidate->getBasicBlocks(BlocksInRegion, BE);
       OS->CE = new (ExtractorAllocator.Allocate())
           CodeExtractor(BE, nullptr, false, nullptr, nullptr, nullptr, false,
-                        false, nullptr, "outlined");
+                        false, nullptr, "outlined", false, true);
       findAddInputsOutputs(M, *OS, NotSame);
       if (!OS->IgnoreRegion)
         OutlinedRegions.push_back(OS);
@@ -2900,7 +2900,7 @@ unsigned IROutliner::doOutline(Module &M) {
       OS->Candidate->getBasicBlocks(BlocksInRegion, BE);
       OS->CE = new (ExtractorAllocator.Allocate())
           CodeExtractor(BE, nullptr, false, nullptr, nullptr, nullptr, false,
-                        false, nullptr, "outlined");
+                        false, nullptr, "outlined", false, true);
       bool FunctionOutlined = extractSection(*OS);
       if (FunctionOutlined) {
         unsigned StartIdx = OS->Candidate->getStartIdx();
diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index 8faf64a60997c..28db49daf2efa 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -264,12 +264,14 @@ CodeExtractor::CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
                              BranchProbabilityInfo *BPI, AssumptionCache *AC,
                              bool AllowVarArgs, bool AllowAlloca,
                              BasicBlock *AllocationBlock, std::string Suffix,
-                             bool ArgsInZeroAddressSpace)
+                             bool ArgsInZeroAddressSpace,
+                             bool VoidReturnWithSingleOutput)
     : DT(DT), AggregateArgs(AggregateArgs || AggregateArgsOpt), BFI(BFI),
       BPI(BPI), AC(AC), AllocationBlock(AllocationBlock),
       AllowVarArgs(AllowVarArgs),
       Blocks(buildExtractionBlockSet(BBs, DT, AllowVarArgs, AllowAlloca)),
-      Suffix(Suffix), ArgsInZeroAddressSpace(ArgsInZeroAddressSpace) {}
+      Suffix(Suffix), ArgsInZeroAddressSpace(ArgsInZeroAddressSpace),
+      VoidReturnWithSingleOutput(VoidReturnWithSingleOutput) {}
 
 /// definedInRegion - Return true if the specified value is defined in the
 /// extracted region.
@@ -662,7 +664,7 @@ bool CodeExtractor::isEligible() const {
 
 void CodeExtractor::findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
                                       const ValueSet &SinkCands,
-                                      bool CollectGlobalInputs) const {
+                                      bool CollectGlobalInputs) {
   for (BasicBlock *BB : Blocks) {
     // If a used value is defined outside the region, it's an input.  If an
     // instruction is used outside the region, it's an output.
@@ -682,6 +684,12 @@ void CodeExtractor::findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
         }
     }
   }
+
+  if (!VoidReturnWithSingleOutput && !AggregateArgs && Outputs.size() == 1 &&
+      getCommonExitBlock(Blocks)) {
+    FuncRetVal = Outputs[0];
+    Outputs.clear();
+  }
 }
 
 /// severSplitPHINodesOfEntry - If a PHI node has multiple inputs from outside
@@ -877,7 +885,7 @@ Function *CodeExtractor::constructFunctionDeclaration(
         M->getContext(), ArgsInZeroAddressSpace ? 0 : DL.getAllocaAddrSpace()));
   }
 
-  Type *RetTy = getSwitchType();
+  Type *RetTy = FuncRetVal ? FuncRetVal->getType() : getSwitchType();
   LLVM_DEBUG({
     dbgs() << "Function type: " << *RetTy << " f(";
     for (Type *i : ParamTy)
@@ -1519,8 +1527,8 @@ CodeExtractor::extractCodeRegion(const CodeExtractorAnalysisCache &CEAC,
       inputs, outputs, StructValues, newFunction, StructTy, oldFunction, ReplIP,
       EntryFreq, LifetimesStart.getArrayRef(), Reloads);
 
-  insertReplacerCall(oldFunction, header, TheCall->getParent(), outputs,
-                     Reloads, ExitWeights);
+  insertReplacerCall(oldFunction, header, TheCall, outputs, Reloads,
+                     ExitWeights);
 
   fixupDebugInfoPostExtraction(*oldFunction, *newFunction, *TheCall, inputs,
                                NewValues);
@@ -1699,14 +1707,17 @@ void CodeExtractor::emitFunctionBody(
     ExitBlockMap[OldTarget] = NewTarget;
 
     Value *brVal = nullptr;
-    Type *RetTy = getSwitchType();
+    Type *RetTy = FuncRetVal ? FuncRetVal->getType() : getSwitchType();
     assert(ExtractedFuncRetVals.size() < 0xffff &&
            "too many exit blocks for switch");
     switch (ExtractedFuncRetVals.size()) {
     case 0:
-    case 1:
       // No value needed.
       break;
+    case 1:
+      if (FuncRetVal)
+        brVal = FuncRetVal;
+      break;
     case 2: // Conditional branch, return a bool
       brVal = ConstantInt::get(RetTy, !SuccNum);
       break;
@@ -2019,13 +2030,14 @@ CallInst *CodeExtractor::emitReplacerCall(
 }
 
 void CodeExtractor::insertReplacerCall(
-    Function *oldFunction, BasicBlock *header, BasicBlock *codeReplacer,
+    Function *oldFunction, BasicBlock *header, CallInst *ReplacerCall,
     const ValueSet &outputs, ArrayRef<Value *> Reloads,
     const DenseMap<BasicBlock *, BlockFrequency> &ExitWeights) {
 
   // Rewrite branches to basic blocks outside of the loop to new dummy blocks
   // within the new function. This must be done before we lose track of which
   // blocks were originally in the code region.
+  BasicBlock *codeReplacer = ReplacerCall->getParent();
   std::vector<User *> Users(header->user_begin(), header->user_end());
   for (auto &U : Users)
     // The BasicBlock which contains the branch is not in the region
@@ -2067,6 +2079,13 @@ void CodeExtractor::insertReplacerCall(
     }
   }
 
+  if (FuncRetVal)
+    for (User *U : FuncRetVal->users()) {
+      Instruction *inst = cast<Instruction>(U);
+      if (inst->getParent()->getParent() == oldFunction)
+        inst->replaceUsesOfWith(FuncRetVal, ReplacerCall);
+    }
+
   // Update the branch weights for the exit block.
   if (BFI && ExtractedFuncRetVals.size() > 1)
     calculateNewCallTerminatorWeights(codeReplacer, ExitWeights, BPI);
diff --git a/llvm/test/Transforms/CodeExtractor/PartialInlineAnd.ll b/llvm/test/Transforms/CodeExtractor/PartialInlineAnd.ll
index 438a598324635..4b3db2a48e572 100644
--- a/llvm/test/Transforms/CodeExtractor/PartialInlineAnd.ll
+++ b/llvm/test/Transforms/CodeExtractor/PartialInlineAnd.ll
@@ -43,7 +43,7 @@ bb:
 ; LIMIT-LABEL: @dummy_caller
 ; LIMIT: br i1
 ; LIMIT-NOT: br
-; LIMIT: call void @bar.1.
+; LIMIT: call i32 @bar.1.
   %tmp = tail call i32 @bar(i32 %arg)
   ret i32 %tmp
 }
diff --git a/llvm/test/Transforms/CodeExtractor/PartialInlineAttributes.ll b/llvm/test/Transforms/CodeExtractor/PartialInlineAttributes.ll
index b085e34d52648..1246c5de6834f 100644
--- a/llvm/test/Transforms/CodeExtractor/PartialInlineAttributes.ll
+++ b/llvm/test/Transforms/CodeExtractor/PartialInlineAttributes.ll
@@ -55,9 +55,9 @@ if.end:
   ret i32 %add
 }
 ; CHECK-LABEL: @caller
-; CHECK: call void @callee_most.2.if.then(i32 %v
+; CHECK: call i32 @callee_most.2.if.then(i32 %v
 ; CHECK: call i32 @callee_noinline(i32 %v)
-; CHECK: call void @callee_writeonly.1.if.then(i32 %v
+; CHECK: call i32 @callee_writeonly.1.if.then(i32 %v
 define i32 @caller(i32 %v) ssp {
 entry:
   %c1 = call i32 @callee_most(i32 %v)
@@ -66,8 +66,8 @@ entry:
   ret i32 %c3
 }
 
-; CHECK: define internal void @callee_writeonly.1.if.then(i32 %v, ptr %sub.out) [[FN_ATTRS0:#[0-9]+]]
-; CHECK: define internal void @callee_most.2.if.then(i32 %v, ptr %sub.out)  [[FN_ATTRS:#[0-9]+]]
+; CHECK: define internal i32 @callee_writeonly.1.if.then(i32 %v) [[FN_ATTRS0:#[0-9]+]]
+; CHECK: define internal i32 @callee_most.2.if.then(i32 %v)  [[FN_ATTRS:#[0-9]+]]
 
 ; attributes to preserve
 attributes #0 = {
diff --git a/llvm/test/Transforms/CodeExtractor/PartialInlineDebug.ll b/llvm/test/Transforms/CodeExtractor/PartialInlineDebug.ll
index ab01bbf20de71..af01ea1421cf5 100644
--- a/llvm/test/Transforms/CodeExtractor/PartialInlineDebug.ll
+++ b/llvm/test/Transforms/CodeExtractor/PartialInlineDebug.ll
@@ -24,7 +24,7 @@ if.end:                                           ; preds = %if.then, %entry
 ; CHECK-LABEL: @caller
 ; CHECK: codeRepl.i:
 ; CHECK-NOT: br label
-; CHECK: call void @callee.2.if.then(i32 %v, ptr %mul.loc.i), !dbg ![[DBG2:[0-9]+]]
+; CHECK: call i32 @callee.2.if.then(i32 %v), !dbg ![[DBG2:[0-9]+]]
 define i32 @caller(i32 %v) !dbg !8 {
 entry:
   %call = call i32 @callee(i32 %v), !dbg !14
@@ -55,17 +55,17 @@ if.end:
 ; CHECK-LABEL: @caller2
 ; CHECK: codeRepl.i:
 ; CHECK-NOT: br label
-; CHECK: call void @callee2.1.if.then(i32 %v, ptr %sub.loc.i), !dbg ![[DBG4:[0-9]+]]
+; CHECK: call i32 @callee2.1.if.then(i32 %v), !dbg ![[DBG4:[0-9]+]]
 define i32 @caller2(i32 %v) !dbg !21 {
 entry:
   %call = call i32 @callee2(i32 %v), !dbg !22
   ret i32 %call
 }
 
-; CHECK-LABEL: define internal void @callee2.1.if.then
+; CHECK-LABEL: define internal i32 @callee2.1.if.then
 ; CHECK: br label %if.then, !dbg ![[DBG5:[0-9]+]]
 
-; CHECK-LABEL: define internal void @callee.2.if.then
+; CHECK-LABEL: define internal i32 @callee.2.if.then
 ; CHECK: br label %if.then, !dbg ![[DBG6:[0-9]+]]
 
 ; CHECK: ![[DBG1]] = !DILocation(line: 10, column: 7,
diff --git a/llvm/test/Transforms/CodeExtractor/PartialInlineInvokeProducesOutVal.ll b/llvm/test/Transforms/CodeExtractor/PartialInlineInvokeProducesOutVal.ll
index 5e0ce20645265..d598174544c5b 100644
--- a/llvm/test/Transforms/CodeExtractor/PartialInlineInvokeProducesOutVal.ll
+++ b/llvm/test/Transforms/CodeExtractor/PartialInlineInvokeProducesOutVal.ll
@@ -24,25 +24,22 @@ bb5:                                              ; preds = %bb4, %bb1, %bb
 
 ; CHECK-LABEL: @dummy_caller
 ; CHECK-LABEL: bb:
-; CHECK-NEXT:  [[CALL26LOC:%.*]] = alloca ptr
 ; CHECK-LABEL: codeRepl.i:
-; CHECK-NEXT:   call void @llvm.lifetime.start.p0(ptr [[CALL26LOC]])
-; CHECK-NEXT:   call void @bar.1.bb1(ptr [[CALL26LOC]])
-; CHECK-NEXT:   %call26.reload.i = load ptr, ptr [[CALL26LOC]]
-; CHECK-NEXT:   call void @llvm.lifetime.end.p0(ptr [[CALL26LOC]])
+; CHECK-NEXT:   call ptr @bar.1.bb1()
 define ptr @dummy_caller(i32 %arg) {
 bb:
   %tmp = tail call ptr @bar(i32 %arg)
   ret ptr %tmp
 }
 
-; CHECK-LABEL: define internal void @bar.1.bb1
+; CHECK-LABEL: define internal ptr @bar.1.bb1
 ; CHECK-LABEL: bb1:
 ; CHECK-NEXT:    %call26 = invoke ptr @invoke_callee()
 ; CHECK-NEXT:            to label %cont unwind label %lpad
 ; CHECK-LABEL: cont:
-; CHECK-NEXT:    store ptr %call26, ptr %call26.out
 ; CHECK-NEXT:    br label %bb5.exitStub
+; CHECK-LABEL: bb5.exitStub:
+; CHECK-NEXT:    ret ptr %call26
 
 ; Function Attrs: nobuiltin
 declare dso_local noalias nonnull ptr @invoke_callee() local_unnamed_addr #1
diff --git a/llvm/test/Transforms/CodeExtractor/PartialInlineOrAnd.ll b/llvm/test/Transforms/CodeExtractor/PartialInlineOrAnd.ll
index 2b357a893419c..6a74c345267d1 100644
--- a/llvm/test/Transforms/CodeExtractor/PartialInlineOrAnd.ll
+++ b/llvm/test/Transforms/CodeExtractor/PartialInlineOrAnd.ll
@@ -54,7 +54,7 @@ bb:
 ; LIMIT3: br i1
 ; LIMIT3: br i1
 ; LIMIT3-NOT: br i1
-; LIMIT3: call void @bar.1.
+; LIMIT3: call i32 @bar.1.
 ; LIMIT2-LABEL: @dummy_caller
 ; LIMIT2-NOT: br i1
 ; LIMIT2: call i32 @bar(
diff --git a/llvm/test/Transforms/CodeExtractor/PartialInlineVarArgsDebug.ll b/llvm/test/Transforms/CodeExtractor/PartialInlineVarArgsDebug.ll
index d0f9e01f950a2..c68729ea231b9 100644
--- a/llvm/test/Transforms/CodeExtractor/PartialInlineVarArgsDebug.ll
+++ b/llvm/test/Transforms/CodeExtractor/PartialInlineVarArgsDebug.ll
@@ -20,14 +20,14 @@ if.end:                                           ; preds = %if.then, %entry
 ; CHECK-LABEL: @caller
 ; CHECK: codeRepl.i:
 ; CHECK-NOT: br label
-; CHECK: call void (i32, ptr, ...) @callee.1.if.then(i32 %v, ptr %mul.loc.i, i32 99), !dbg ![[DBG2:[0-9]+]]
+; CHECK: call i32 (i32, ...) @callee.1.if.then(i32 %v, i32 99), !dbg ![[DBG2:[0-9]+]]
 define i32 @caller(i32 %v) !dbg !8 {
 entry:
   %call = call i32 (i32, ...) @callee(i32 %v, i32 99), !dbg !14
   ret i32 %call, !dbg !15
 }
 
-; CHECK-LABEL: define internal void @callee.1.if.then
+; CHECK-LABEL: define internal i32  @callee.1.if.then
 ; CHECK: br label %if.then, !dbg ![[DBG3:[0-9]+]]
 
 ; CHECK: ![[DBG1]] = !DILocation(line: 10, column: 7,
diff --git a/llvm/test/Transforms/HotColdSplit/apply-penalty-for-outputs.ll b/llvm/test/Transforms/HotColdSplit/apply-penalty-for-outputs.ll
index 96d73088e7874..d4cfccf2bff15 100644
--- a/llvm/test/Transforms/HotColdSplit/apply-penalty-for-outputs.ll
+++ b/llvm/test/Transforms/HotColdSplit/apply-penalty-for-outputs.ll
@@ -11,9 +11,9 @@ entry:
 
 cold:
   ; CHECK: Applying penalty for splitting: 2
-  ; CHECK-NEXT: Applying penalty for: 1 params
-  ; CHECK-NEXT: Applying penalty for: 1 outputs/split phis
-  ; CHECK-NEXT: penalty = 7
+  ; CHECK-NEXT: Applying penalty for: 0 params
+  ; CHECK-NEXT: Applying penalty for: 0 outputs/split phis
+  ; CHECK-NEXT: penalty = 2
   %local = load i32, ptr @g
   call void @sink()
   br label %exit
diff --git a/llvm/test/Transforms/HotColdSplit/duplicate-phi-preds-crash.ll b/llvm/test/Transforms/HotColdSplit/duplicate-phi-preds-crash.ll
index ec772ce181a12..de6d60c825a40 100644
--- a/llvm/test/Transforms/HotColdSplit/duplicate-phi-preds-crash.ll
+++ b/llvm/test/Transforms/HotColdSplit/duplicate-phi-preds-crash.ll
@@ -15,9 +15,9 @@ declare void @sink() cold
 ; CHECK: call {{.*}}@sideeffect(
 ; CHECK: call {{.*}}@realloc(
 ; CHECK-LABEL: codeRepl:
-; CHECK: call {{.*}}@realloc2.cold.1(i64 %size, ptr %ptr, ptr %retval.0.ce.loc)
+; CHECK: %2 = call {{.*}}@realloc2.cold.1(i64 %size, ptr %ptr)
 ; CHECK-LABEL: cleanup:
-; CHECK-NEXT: phi ptr [ null, %if.then ], [ %call, %if.end ], [ %retval.0.ce.reload, %codeRepl ]
+; CHECK-NEXT: phi ptr [ null, %if.then ], [ %call, %if.end ], [ %2, %codeRepl ]
 define ptr @realloc2(ptr %ptr, i64 %size) {
 entry:
   %0 = add i64 %size, -1
diff --git a/llvm/test/Transforms/HotColdSplit/phi-with-distinct-outlined-values.ll b/llvm/test/Transforms/HotColdSplit/phi-with-distinct-outlined-values.ll
index 7c07cd35a527c..d353adef057af 100644
--- a/llvm/test/Transforms/HotColdSplit/phi-with-distinct-outlined-values.ll
+++ b/llvm/test/Transforms/HotColdSplit/phi-with-distinct-outlined-values.ll
@@ -4,12 +4,13 @@ target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-apple-macosx10.14.0"
 
 ; CHECK-LABEL: define {{.*}}@foo(
-; CHECK: phi i32 [ 0, %entry ], [ %p.ce.reload, %codeRepl ]
+; CHECK: %0 = call {{.*}}@foo.cold.1(i1 %arg) 
+; CHECK: phi i32 [ 0, %entry ], [ %0, %codeRepl ]
 
 ; CHECK-LABEL: define {{.*}}@foo.cold.1(
 ; CHECK: call {{.*}}@sink
 ; CHECK: %p.ce = phi i32 [ 1, %coldbb ], [ 3, %coldbb2 ]
-; CHECK-NEXT: store i32 %p.ce, ptr %p.ce.out
+; CHECK: ret i32 %p.ce
 
 define void @foo(i32 %cond, i1 %arg) {
 entry:
diff --git a/llvm/test/Transforms/HotColdSplit/split-phis-in-exit-blocks.ll b/llvm/test/Transforms/HotColdSplit/split-phis-in-exit-blocks.ll
index 915b3393a53b4..8caea59090803 100644
--- a/llvm/test/Transforms/HotColdSplit/split-phis-in-exit-blocks.ll
+++ b/llvm/test/Transforms/HotColdSplit/split-phis-in-exit-blocks.ll
@@ -5,21 +5,17 @@ target triple = "x86_64-apple-macosx10.14.0"
 
 ; CHECK-LABEL: define {{.*}}@pluto(
 ; CHECK-NEXT: bb:
-; CHECK-NEXT:  %tmp8.ce.loc = alloca i1
 ; CHECK-NEXT:  switch i8 undef, label %codeRepl [
 ; CHECK-NEXT:    i8 0, label %bb7
 ; CHECK-NEXT:    i8 1, label %bb7
 ; CHECK-NEXT:  ]
 ;
 ; CHECK:  codeRepl:
-; CHECK-NEXT:    lifetime.start
-; CHECK-NEXT:    call void @pluto.cold.1(ptr %tmp8.ce.loc)
-; CHECK-NEXT:    %tmp8.ce.reload = load i1, ptr %tmp8.ce.loc
-; CHECK-NEXT:    lifetime.end
+; CHECK-NEXT:    %0 = call i1 @pluto.cold.1()
 ; CHECK-NEXT:    br label %bb7
 ;
 ; CHECK:  bb7:
-; CHECK:    %tmp8 = phi i1 [ true, %bb ], [ true, %bb ], [ %tmp8.ce.reload, %codeRepl ]
+; CHECK:    %tmp8 = phi i1 [ true, %bb ], [ true, %bb ], [ %0, %codeRepl ]
 ; CHECK:    ret void
 
 ; CHECK-LABEL: define {{.*}}@pluto.cold.1(
diff --git a/llvm/test/Transforms/HotColdSplit/succ-block-with-self-edge.ll b/llvm/test/Transforms/HotColdSplit/succ-block-with-self-edge.ll
index 59efa3d54e249..e528cb7f7965d 100644
--- a/llvm/test/Transforms/HotColdSplit/succ-block-with-self-edge.ll
+++ b/llvm/test/Transforms/HotColdSplit/succ-block-with-self-edge.ll
@@ -4,9 +4,9 @@ target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-apple-macosx10.14.0"
 
 ; CHECK-LABEL: define {{.*}}@exit_block_with_same_incoming_vals
-; CHECK: call {{.*}}@exit_block_with_same_incoming_vals.cold.1(
+; CHECK: %0 = call {{.*}}@exit_block_with_same_incoming_vals.cold.1(
 ; CHECK-NOT: br i1 %arg
-; CHECK: phi i32 [ 0, %entry ], [ %p.ce.reload, %codeRepl ]
+; CHECK: phi i32 [ 0, %entry ], [ %0, %codeRepl ]
 define void @exit_block_with_same_incoming_vals(i32 %cond, i1 %arg) {
 entry:
   %tobool = icmp eq i32 %cond, 0
@@ -27,9 +27,9 @@ if.end:
 }
 
 ; CHECK-LABEL: define {{.*}}@exit_block_with_distinct_incoming_vals
-; CHECK: call {{.*}}@exit_block_with_distinct_incoming_vals.cold.1(
+; CHECK: %0 = call {{.*}}@exit_block_with_distinct_incoming_vals.cold.1(
 ; CHECK-NOT: br i1 %arg
-; CHECK: phi i32 [ 0, %entry ], [ %p.ce.reload, %codeRepl ]
+; CHECK: phi i32 [ 0, %entry ], [ %0, %codeRepl ]
 define void @exit_block_with_distinct_incoming_vals(i32 %cond, i1 %arg) {
 entry:
   %tobool = icmp eq i32 %cond, 0
diff --git a/llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp b/llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp
index 90f06204ec9b3..1e4b530a901d4 100644
--- a/llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp
+++ b/llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp
@@ -69,7 +69,8 @@ TEST(CodeExtractor, ExitStub) {
                                            getBlockByName(Func, "body1"),
                                            getBlockByName(Func, "body2") };
 
-  CodeExtractor CE(Candidates);
+  CodeExtractor CE(Candidates, nullptr, false, nullptr, nullptr, nullptr, false,
+                   false, nullptr, "", false, true);
   EXPECT_TRUE(CE.isEligible());
 
   CodeExtractorAnalysisCache CEAC(*Func);
@@ -117,7 +118,8 @@ TEST(CodeExtractor, InputOutputMonitoring) {
                                           getBlockByName(Func, "body1"),
                                           getBlockByName(Func, "body2")};
 
-  CodeExtractor CE(Candidates);
+  CodeExtractor CE(Candidates, nullptr, false, nullptr, nullptr, nullptr, false,
+                   false, nullptr, "", false, true);
   EXPECT_TRUE(CE.isEligible());
 
   CodeExtractorAnalysisCache CEAC(*Func);
@@ -145,6 +147,70 @@ TEST(CodeExtractor, InputOutputMonitoring) {
   EXPECT_FALSE(verifyFunction(*Func));
 }
 
+TEST(CodeExtractor, InputOutputReturnMonitoring) {
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M(parseAssemblyString(R"invalid(
+    define i32 @foo(i32 %x, i32 %y, i32 %z) {
+    header:
+      %0 = icmp ugt i32 %x, %y
+      br i1 %0, label %body1, label %body2
+
+    body1:
+      %1 = add i32 %z, 2
+      br label %notExtracted
+
+    body2:
+      %2 = mul i32 %z, 7
+      br label %notExtracted
+
+    notExtracted:
+      %3 = phi i32 [ %1, %body1 ], [ %2, %body2 ]
+      %4 = add i32 %3, %x
+      ret i32 %4
+    }
+  )invalid",
+                                                Err, Ctx));
+
+  Function *Func = M->getFunction("foo");
+  SmallVector<BasicBlock *, 3> Candidates{getBlockByName(Func, "header"),
+                                          getBlockByName(Func, "body1"),
+                                          getBlockByName(Func, "body2")};
+
+  CodeExtractor CE(Candidates, nullptr, false, nullptr, nullptr, nullptr, false,
+                   false, nullptr, "", false, false);
+  EXPECT_TRUE(CE.isEligible());
+
+  CodeExtractorAnalysisCache CEAC(*Func);
+  SetVector<Value *> Inputs, Outputs;
+  Function *Outlined = CE.extractCodeRegion(CEAC, Inputs, Outputs);
+  EXPECT_TRUE(Outlined);
+
+  EXPECT_EQ(Inputs.size(), 3u);
+  EXPECT_EQ(Inputs[0], Func->getArg(2));
+  EXPECT_EQ(Inputs[1], Func->getArg(0));
+  EXPECT_EQ(Inputs[2], Func->getArg(1));
+  EXPECT_EQ(Outputs.size(), 0u);
+  BasicBlock *Exit = getBlockByName(Func, "notExtracted");
+  BasicBlock *ExitSplit = getBlockByName(Outlined, "notExtracted.split");
+  // Ensure that PHI in exit block has only one incoming value (from code
+  // replacer block).
+  EXPECT_TRUE(Exit && cast<PHINode>(Exit->front()).getNumIncomingValues() == 1);
+  // Ensure that there is a PHI in outlined function with 2 incoming values.
+  EXPECT_TRUE(ExitSplit &&
+              cast<PHINode>(ExitSplit->front()).getNumIncomingValues() == 2);
+
+  BasicBlock *ExitStub = getBlockByName(Outlined, "notExtracted.exitStub");
+  Instruction *ExitTerm = ExitStub->getTerminator();
+  ReturnInst *ExitReturn = dyn_cast<ReturnInst>(ExitTerm);
+  EXPECT_TRUE(ExitReturn);
+  Value *RetVal = ExitReturn->getReturnValue();
+  EXPECT_TRUE(ExitReturn);
+
+  EXPECT_FALSE(verifyFunction(*Outlined));
+  EXPECT_FALSE(verifyFunction(*Func));
+}
+
 TEST(CodeExtractor, ExitBlockOrderingPhis) {
   LLVMContext Ctx;
   SMDiagnostic Err;
@@ -177,7 +243,8 @@ TEST(CodeExtractor, ExitBlockOrderingPhis) {
                                            getBlockByName(Func, "test1"),
                                            getBlockByName(Func, "test") };
 
-  CodeExtractor CE(Candidates);
+  CodeExtractor CE(Candidates, nullptr, false, nullptr, nullptr, nullptr, false,
+                   false, nullptr, "", false, true);
   EXPECT_TRUE(CE.isEligible());
 
   CodeExtractorAnalysisCache CEAC(*Func);
@@ -234,7 +301,8 @@ TEST(CodeExtractor, ExitBlockOrdering) {
                                            getBlockByName(Func, "test1"),
                                            getBlockByName(Func, "test") };
 
-  CodeExtractor CE(Candidates);
+  CodeExtractor CE(Candidates, nullptr, false, nullptr, nullptr, nullptr, false,
+                   false, nullptr, "", false, true);
   EXPECT_TRUE(CE.isEligible());
 
   CodeExtractorAnalysisCache CEAC(*Func);
@@ -293,7 +361,8 @@ TEST(CodeExtractor, ExitPHIOnePredFromRegion) {
     getBlockByName(Func, "extracted2")
   };
 
-  CodeExtractor CE(ExtractedBlocks);
+  CodeExtractor CE(ExtractedBlocks, nullptr, false, nullptr, nullptr, nullptr,
+                   false, false, nullptr, "", false, true);
   EXPECT_TRUE(CE.isEligible());
 
   CodeExtractorAnalysisCache CEAC(*Func);
@@ -353,7 +422,7 @@ TEST(CodeExtractor, StoreOutputInvokeResultAfterEHPad) {
     }
   )invalid", Err, Ctx));
 
-	if (!M) {
+  if (!M) {
     Err.print("unit", errs());
     exit(1);
   }
@@ -368,7 +437,8 @@ TEST(CodeExtractor, StoreOutputInvokeResultAfterEHPad) {
     getBlockByName(Func, "lpad2")
   };
 
-  CodeExtractor CE(ExtractedBlocks);
+  CodeExtractor CE(ExtractedBlocks, nullptr, false, nullptr, nullptr, nullptr,
+                   false, false, nullptr, "", false, true);
   EXPECT_TRUE(CE.isEligible());
 
   CodeExtractorAnalysisCache CEAC(*Func);
@@ -403,7 +473,8 @@ TEST(CodeExtractor, StoreOutputInvokeResultInExitStub) {
   SmallVector<BasicBlock *, 1> Blocks{ getBlockByName(Func, "entry"),
                                        getBlockByName(Func, "lpad") };
 
-  CodeExtractor CE(Blocks);
+  CodeExtractor CE(Blocks, nullptr, false, nullptr, nullptr, nullptr, false,
+                   false, nullptr, "", false, true);
   EXPECT_TRUE(CE.isEligible());
 
   CodeExtractorAnalysisCache CEAC(*Func);
@@ -455,7 +526,8 @@ TEST(CodeExtractor, ExtractAndInvalidateAssumptionCache) {
   Function *Func = M->getFunction("test");
   SmallVector<BasicBlock *, 1> Blocks{ getBlockByName(Func, "if.else") };
   AssumptionCache AC(*Func);
-  CodeExtractor CE(Blocks, nullptr, false, nullptr, nullptr, &AC);
+  CodeExtractor CE(Blocks, nullptr, false, nullptr, nullptr, &AC, false, false,
+                   nullptr, "", false, true);
   EXPECT_TRUE(CE.isEligible());
 
   CodeExtractorAnalysisCache CEAC(*Func);
@@ -499,7 +571,8 @@ TEST(CodeExtractor, RemoveBitcastUsesFromOuterLifetimeMarkers) {
   Function *Func = M->getFunction("foo");
   SmallVector<BasicBlock *, 1> Blocks{getBlockByName(Func, "extract")};
 
-  CodeExtractor CE(Blocks);
+  CodeExtractor CE(Blocks, nullptr, false, nullptr, nullptr, nullptr, false,
+                   false, nullptr, "", false, true);
   EXPECT_TRUE(CE.isEligible());
 
   CodeExtractorAnalysisCache CEAC(*Func);
@@ -554,7 +627,8 @@ TEST(CodeExtractor, PartialAggregateArgs) {
 
   // Create the CodeExtractor with arguments aggregation enabled.
   CodeExtractor CE(Blocks, /* DominatorTree */ nullptr,
-                   /* AggregateArgs */ true);
+                   /* AggregateArgs */ true, nullptr, nullptr, nullptr, false,
+                   false, nullptr, "", false, true);
   EXPECT_TRUE(CE.isEligible());
 
   CodeExtractorAnalysisCache CEAC(*Func);
@@ -602,7 +676,7 @@ TEST(CodeExtractor, AllocaBlock) {
 
   BasicBlock *AllocaBlock = getBlockByName(Func, "allocas");
   CodeExtractor CE(Candidates, nullptr, true, nullptr, nullptr, nullptr, false,
-                   false, AllocaBlock);
+                   false, AllocaBlock, "", false, true);
   CE.excludeArgFromAggregate(Func->getArg(0));
   CE.excludeArgFromAggregate(getInstByName(Func, "w"));
   EXPECT_TRUE(CE.isEligible());
@@ -656,7 +730,8 @@ TEST(CodeExtractor, PartialAggregateArgs2) {
 
   // Create the CodeExtractor with arguments aggregation enabled.
   CodeExtractor CE(Blocks, /* DominatorTree */ nullptr,
-                   /* AggregateArgs */ true);
+                   /* AggregateArgs */ true, nullptr, nullptr, nullptr, false,
+                   false, nullptr, "", false, true);
   EXPECT_TRUE(CE.isEligible());
 
   CodeExtractorAnalysisCache CEAC(*Func);
@@ -713,7 +788,8 @@ TEST(CodeExtractor, OpenMPAggregateArgs) {
                    /* AllowAlloca */ true,
                    /* AllocaBlock*/ &Func->getEntryBlock(),
                    /* Suffix */ ".outlined",
-                   /* ArgsInZeroAddressSpace */ true);
+                   /* ArgsInZeroAddressSpace */ true,
+                   /* VoidReturnWithSingleOutput */ true);
 
   EXPECT_TRUE(CE.isEligible());
 
@@ -779,7 +855,9 @@ TEST(CodeExtractor, ArgsDebugInfo) {
   SmallVector<BasicBlock *, 1> Blocks{getBlockByName(Func, "extract")};
 
   auto TestExtracted = [&](bool AggregateArgs) {
-    CodeExtractor CE(Blocks, /* DominatorTree */ nullptr, AggregateArgs);
+    CodeExtractor CE(Blocks, /* DominatorTree */ nullptr, AggregateArgs,
+                     nullptr, nullptr, nullptr, false, false, nullptr, "",
+                     false, true);
     EXPECT_TRUE(CE.isEligible());
     CodeExtractorAnalysisCache CEAC(*Func);
     SetVector<Value *> Inputs, Outputs, SinkingCands, HoistingCands;



More information about the llvm-commits mailing list