[llvm] [SimplifyCFG] Update profile metadata regardless of weight count (PR #190982)

[Aiden Grossman via llvm-commits]

https://github.com/boomanaiden154 updated https://github.com/llvm/llvm-project/pull/190982

>From da0682b3e5ff789e616e92ec429fbff7de7c0f7d Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Wed, 8 Apr 2026 14:41:52 +0000
Subject: [PATCH 1/2] [SimplifyCFG] Update profile metadata regardless of
 weight count

Otherwise we miss updating in cases where we should be updating which
causes assertion failures later due to the fact that the number of cases
no longer matches the number of branch weights.

Fixes #190901.
---
 llvm/include/llvm/IR/Instructions.h           |  2 +-
 .../SimplifyCFG/profile-update-unreachable.ll | 39 +++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/SimplifyCFG/profile-update-unreachable.ll

diff --git a/llvm/include/llvm/IR/Instructions.h b/llvm/include/llvm/IR/Instructions.h
index bdc6d5bd2f5e5..9f347a39e6671 100644
--- a/llvm/include/llvm/IR/Instructions.h
+++ b/llvm/include/llvm/IR/Instructions.h
@@ -3699,7 +3699,7 @@ class SwitchInstProfUpdateWrapper {
   SwitchInstProfUpdateWrapper(SwitchInst &SI) : SI(SI) { init(); }
 
   ~SwitchInstProfUpdateWrapper() {
-    if (Changed && Weights.has_value() && Weights->size() >= 2)
+    if (Changed && Weights.has_value())
       setBranchWeights(SI, Weights.value(), /*IsExpected=*/false);
   }
 
diff --git a/llvm/test/Transforms/SimplifyCFG/profile-update-unreachable.ll b/llvm/test/Transforms/SimplifyCFG/profile-update-unreachable.ll
new file mode 100644
index 0000000000000..fcafb4a716c87
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/profile-update-unreachable.ll
@@ -0,0 +1,39 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -passes="simplifycfg" -S | FileCheck %s
+
+define i32 @f1(i1 %cond1, i8 %cond2) {
+; CHECK-LABEL: define i32 @f1(
+; CHECK-SAME: i1 [[COND1:%.*]], i8 [[COND2:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND1]])
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  br i1 %cond1, label %swbb1, label %swbb2
+
+swbb1.case1:
+  ret i32 0
+
+swbb2.case1:
+  unreachable
+
+sw.case2:
+  unreachable
+
+sw.default:
+  unreachable
+
+swbb1:
+  switch i8 %cond2, label %sw.default [
+  i8 0, label %swbb1.case1
+  i8 1, label %sw.case2
+  ]
+
+swbb2:
+  switch i8 %cond2, label %sw.default [
+  i8 0, label %swbb2.case1
+  i8 1, label %sw.case2
+  ], !prof !0
+}
+
+!0 = !{!"branch_weights", i32 3, i32 5, i32 7}

>From 295e03934a7a50c764c9a4f48d066a33490f9863 Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Wed, 8 Apr 2026 15:27:52 +0000
Subject: [PATCH 2/2] feedback

---
 llvm/include/llvm/IR/Instructions.h | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/IR/Instructions.h b/llvm/include/llvm/IR/Instructions.h
index 9f347a39e6671..d0d81971e97ac 100644
--- a/llvm/include/llvm/IR/Instructions.h
+++ b/llvm/include/llvm/IR/Instructions.h
@@ -3699,8 +3699,17 @@ class SwitchInstProfUpdateWrapper {
   SwitchInstProfUpdateWrapper(SwitchInst &SI) : SI(SI) { init(); }
 
   ~SwitchInstProfUpdateWrapper() {
-    if (Changed && Weights.has_value())
-      setBranchWeights(SI, Weights.value(), /*IsExpected=*/false);
+    if (Changed && Weights.has_value()) {
+      if (Weights->size() >= 2) {
+        setBranchWeights(SI, Weights.value(), /*IsExpected=*/false);
+        return;
+      }
+      // In some cases while simplifying switch instructions, we end up with
+      // degenerate switch instructions (e.g., only contain the default case).
+      // We drop profile metadata in such cases given it does not convey
+      // anything.
+      SI.setMetadata(LLVMContext::MD_prof, nullptr);
+    }
   }
 
   /// Delegate the call to the underlying SwitchInst::removeCase() and remove