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

scui-ibm · web-flow · commit 66f0da73886f · 2026-04-20T16:55:50.000-04:00
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.
diff --git 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 = true);
 
   /// 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/HotColdSplitting.cpp b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
@@ -721,7 +721,9 @@ bool HotColdSplitting::outlineColdRegions(Function &F, bool HasProfileSummary) {
             SubRegion, &*DT, /* AggregateArgs */ false, /* BFI */ nullptr,
             /* BPI */ nullptr, AC, /* AllowVarArgs */ false,
             /* AllowAlloca */ false, /* AllocaBlock */ nullptr,
-            /* Suffix */ "cold." + std::to_string(OutlinedFunctionID));
+            /* Suffix */ "cold." + std::to_string(OutlinedFunctionID),
+            /* ArgsInZeroAddressSpace */ false,
+            /* VoidReturnWithSingleOutput */ false);
 
         if (CE.isEligible() && isSplittingBeneficial(CE, SubRegion, TTI) &&
             // If this outlining region intersects with another, drop the new
diff --git a/llvm/lib/Transforms/IPO/PartialInlining.cpp b/llvm/lib/Transforms/IPO/PartialInlining.cpp
@@ -1104,7 +1104,10 @@ bool PartialInlinerImpl::FunctionCloner::doMultiRegionFunctionOutlining() {
     CodeExtractor CE(RegionInfo.Region, &DT, /*AggregateArgs*/ false,
                      ClonedFuncBFI.get(), &BPI,
                      LookupAC(*RegionInfo.EntryBlock->getParent()),
-                     /* AllowVarargs */ false);
+                     /* AllowVarargs */ false, /* AllowAlloca */ false,
+                     /* AllocaBlock */ nullptr, /* Suffix */ "",
+                     /* ArgsInZeroAddressSpace */ false,
+                     /* VoidReturnWithSingleOutput */ false);
 
     CE.findInputsOutputs(Inputs, Outputs, Sinks);
 
@@ -1185,7 +1188,10 @@ PartialInlinerImpl::FunctionCloner::doSingleRegionFunctionOutlining() {
   Function *OutlinedFunc =
       CodeExtractor(ToExtract, &DT, /*AggregateArgs*/ false,
                     ClonedFuncBFI.get(), &BPI, LookupAC(*ClonedFunc),
-                    /* AllowVarargs */ true)
+                    /* AllowVarargs */ true, /* AllowAlloca */ false,
+                    /* AllocaBlock */ nullptr, /* Suffix */ "",
+                    /* ArgsInZeroAddressSpace */ false,
+                    /* VoidReturnWithSingleOutput */ false)
           .extractCodeRegion(CEAC);
 
   if (OutlinedFunc) {
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
diff --git a/llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp b/llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp

Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,7 @@ bb:`
`43`	`43`	`; LIMIT-LABEL: @dummy_caller`
`44`	`44`	`; LIMIT: br i1`
`45`	`45`	`; LIMIT-NOT: br`
`46`		`-; LIMIT: call void @bar.1.`
	`46`	`+; LIMIT: call i32 @bar.1.`
`47`	`47`	`%tmp = tail call i32 @bar(i32 %arg)`
`48`	`48`	`ret i32 %tmp`
`49`	`49`	`}`