[llvm] [CHR] Fix crash when marking merged condition unknown (PR #173902)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Peter Waller (peterwaller-arm)

<details>
<summary>Changes</summary>

CHR builds the merged hot-path predicate with IRBuilder::CreateLogicalAnd. That helper is implemented as a select and can constant-fold to a non- Instruction (e.g. i1 true). The pass then attempted to mark the merged condition as having explicitly unknown branch weights when profile data is present, but it unconditionally did cast<Instruction>(MergedCondition), which can crash in release builds.

Guard the metadata update with dyn_cast<Instruction> and pass the containing Function explicitly to avoid calling Instruction::getFunction when the value is not attached yet.

Add a regression test that exercises the constant-folding case.

Crashing stack:

```
  2.      Running pass "chr" on function "repro_crash"
  #<!-- -->0 0x0000000003be00a4 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (bin/opt+0x3be00a4)
  #<!-- -->1 0x0000000003bdd9e8 llvm::sys::RunSignalHandlers() (bin/opt+0x3bdd9e8)
  #<!-- -->2 0x0000000003be1300 SignalHandler(int, siginfo_t*, void*) Signals.cpp:0:0
  #<!-- -->3 0x0000ffffa8e1d840 (linux-vdso.so.1+0x840)
  #<!-- -->4 0x0000000003c815e0 llvm::Instruction::getFunction() const (bin/opt+0x3c815e0)
  #<!-- -->5 0x0000000003dcd35c llvm::setExplicitlyUnknownBranchWeightsIfProfiled(llvm::Instruction&, llvm::StringRef, llvm::Function const*) (bin/opt+0x3dcd35c)
  #<!-- -->6 0x0000000004fb3670 (anonymous namespace)::CHR::addToMergedCondition(bool, llvm::Value*, llvm::Instruction*, (anonymous namespace)::CHRScope*, llvm::IRBuilder<llvm::ConstantFolder, llvm::IRBuilderDefaultInserter>&, llvm::Value*&) ControlHeightReduction.cpp:0:0
  #<!-- -->7 0x0000000004fa7d88 (anonymous namespace)::CHR::run() ControlHeightReduction.cpp:0:0
  #<!-- -->8 0x0000000004fa3618 llvm::ControlHeightReductionPass::run(llvm::Function&, llvm::AnalysisManager<llvm::Function>&) (bin/opt+0x4fa3618)
```

Tests: opt < llvm/test/Transforms/PGOProfile/chr-unknown-profdata-crash.ll -passes='require<profile-summary>,function(chr)' -force-chr -chr-merge-threshold=1 -disable-output

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp (+2-2) 
- (added) llvm/test/Transforms/PGOProfile/chr-unknown-profdata-crash.ll (+47) 


``````````diff
diff --git a/llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp b/llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp
index 726d94b27a7f2..c7b941319f8b9 100644
--- a/llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp
+++ b/llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp
@@ -1992,8 +1992,8 @@ void CHR::addToMergedCondition(bool IsTrueBiased, Value *Cond,
 
   // Use logical and to avoid propagating poison from later conditions.
   MergedCondition = IRB.CreateLogicalAnd(MergedCondition, Cond);
-  setExplicitlyUnknownBranchWeightsIfProfiled(
-      *cast<Instruction>(MergedCondition), DEBUG_TYPE);
+  if (auto *MergedInst = dyn_cast<Instruction>(MergedCondition))
+    setExplicitlyUnknownBranchWeightsIfProfiled(*MergedInst, DEBUG_TYPE, &F);
 }
 
 void CHR::transformScopes(SmallVectorImpl<CHRScope *> &CHRScopes) {
diff --git a/llvm/test/Transforms/PGOProfile/chr-unknown-profdata-crash.ll b/llvm/test/Transforms/PGOProfile/chr-unknown-profdata-crash.ll
new file mode 100644
index 0000000000000..a28dc0c5287c3
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/chr-unknown-profdata-crash.ll
@@ -0,0 +1,47 @@
+; RUN: opt < %s -passes='require<profile-summary>,function(chr)' -force-chr -chr-merge-threshold=1 -disable-output
+
+declare void @foo()
+declare void @bar()
+
+; Regression test for a crash in CHR when setting unknown profdata on the
+; merged condition. IRBuilder::CreateLogicalAnd is implemented as a select and
+; can constant-fold to a non-Instruction value (e.g. `i1 true`). The buggy code
+; assumed it always produced an Instruction and did `cast<Instruction>(V)`,
+; which can segfault in release builds.
+define void @repro_crash() !prof !14 {
+entry:
+  br i1 true, label %then1, label %cont1, !prof !15
+
+then1:
+  call void @foo()
+  br label %cont1
+
+cont1:
+  br i1 true, label %then2, label %exit, !prof !15
+
+then2:
+  call void @bar()
+  br label %exit
+
+exit:
+  ret void
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 1, !"ProfileSummary", !1}
+!1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+!2 = !{!"ProfileFormat", !"InstrProf"}
+!3 = !{!"TotalCount", i64 10000}
+!4 = !{!"MaxCount", i64 10}
+!5 = !{!"MaxInternalCount", i64 1}
+!6 = !{!"MaxFunctionCount", i64 1000}
+!7 = !{!"NumCounts", i64 1}
+!8 = !{!"NumFunctions", i64 1}
+!9 = !{!"DetailedSummary", !10}
+!10 = !{!11, !12, !13}
+!11 = !{i32 10000, i64 100, i32 1}
+!12 = !{i32 999000, i64 100, i32 1}
+!13 = !{i32 999999, i64 1, i32 2}
+
+!14 = !{!"function_entry_count", i64 100}
+!15 = !{!"branch_weights", i32 1000000, i32 1}

``````````

</details>


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