[llvm] [CodeExtractor] Reset stale FuncRetVal in findInputsOutputs (PR #195484)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Steven Wu (cachemeifyoucan)

<details>
<summary>Changes</summary>

HotColdSplitting calls findInputsOutputs once for cost analysis before
extractCodeRegion runs normalizeCFGForExtraction. The first call may set
FuncRetVal to a single-output value which severSplitPHINodesOfEntry then
moves back into the parent function. The second findInputsOutputs call
correctly finds zero outputs but leaves FuncRetVal pointing at the
parent's instruction, so the cold function emits a `ret` referencing a
value in another function and fails verification.

Reset FuncRetVal at the start of the selection block so each call
recomputes it from the current CFG.

rdar://176075761


---
Full diff: https://github.com/llvm/llvm-project/pull/195484.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Utils/CodeExtractor.cpp (+3) 
- (added) llvm/test/Transforms/HotColdSplit/stale-funcretval-after-sever.ll (+32) 


``````````diff
diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index d2215542ff785..949aae3329589 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -707,6 +707,9 @@ void CodeExtractor::findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
     }
   }
 
+  // Reset stale state from any prior call in HotColdSplitting; the CFG may
+  // have changed since.
+  FuncRetVal = nullptr;
   if (!VoidReturnWithSingleOutput && !AggregateArgs && Outputs.size() == 1 &&
       getCommonExitBlock(Blocks)) {
     FuncRetVal = Outputs[0];
diff --git a/llvm/test/Transforms/HotColdSplit/stale-funcretval-after-sever.ll b/llvm/test/Transforms/HotColdSplit/stale-funcretval-after-sever.ll
new file mode 100644
index 0000000000000..cac54ac7b9dc0
--- /dev/null
+++ b/llvm/test/Transforms/HotColdSplit/stale-funcretval-after-sever.ll
@@ -0,0 +1,32 @@
+; RUN: opt -S -passes=hotcoldsplit -hotcoldsplit-threshold=0 < %s | FileCheck %s
+
+target datalayout = "e-m:o-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-n32:64-S128-Fn32"
+target triple = "arm64-apple-macosx27.0.0"
+
+declare void @sink_call(ptr, i32, ptr)
+
+; CHECK-LABEL: define ptr @test_stale_funcretval(
+; CHECK: call void @test_stale_funcretval.cold.1()
+; CHECK-LABEL: end:
+; CHECK-NEXT: phi ptr [ %ph, %codeRepl ], [ null, %then ]
+
+; CHECK-LABEL: define internal void @test_stale_funcretval.cold.1()
+; CHECK: ret void
+define ptr @test_stale_funcretval() {
+entry:
+  br i1 false, label %sink, label %then
+
+then:                                             ; preds = %entry
+  br i1 false, label %end, label %sink, !prof !0
+
+sink:                                             ; preds = %then, %entry
+  %ph = phi ptr [ null, %then ], [ null, %entry ]
+  tail call void @sink_call(ptr null, i32 0, ptr null)
+  br label %end
+
+end:                                              ; preds = %sink, %then
+  %val = phi ptr [ %ph, %sink ], [ null, %then ]
+  ret ptr %val
+}
+
+!0 = !{!"branch_weights", !"expected", i32 2000, i32 1}

``````````

</details>


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