[alokkrsharma]

Reassociating (A || B) || C and (A && B) && C leaves a select that still tests A. Copy that select's existing weights instead of marking the new one unknown.

[llvmorg-github-actions]

@llvm/pr-subscribers-llvm-transforms

Author: Alok Kumar Sharma (alokkrsharma)

Changes

Reassociating (A || B) || C and (A && B) && C leaves a select that still tests A. Copy that select's existing weights instead of marking the new one unknown.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp (+20-8)
• (added) llvm/test/Transforms/InstCombine/select-logical-reassoc.ll (+56)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 5b2b953953d58..77ab6d2397a68 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -3892,10 +3892,16 @@ Instruction *InstCombinerImpl::foldSelectOfBools(SelectInst &SI) {
     if (match(CondVal, m_OneUse(m_Select(m_Value(A), m_One(), m_Value(B)))) &&
         impliesPoisonOrCond(FalseVal, B, /*Expected=*/false, SQ)) {
       // (A || B) || C --> A || (B | C)
-      Value *LOr = Builder.CreateLogicalOr(A, Builder.CreateOr(B, FalseVal));
-      if (auto *I = dyn_cast<Instruction>(LOr)) {
-        setExplicitlyUnknownBranchWeightsIfProfiled(*I, DEBUG_TYPE);
-      }
+      // CondVal is select A, true, B, so its weights describe A.
+      Instruction *MDFrom = nullptr;
+      if (!ProfcheckDisableMetadataFixes &&
+          cast<Instruction>(CondVal)->getMetadata(LLVMContext::MD_prof))
+        MDFrom = cast<Instruction>(CondVal);
+      Value *LOr =
+          Builder.CreateLogicalOr(A, Builder.CreateOr(B, FalseVal), "", MDFrom);
+      if (!ProfcheckDisableMetadataFixes && !MDFrom)
+        if (auto *I = dyn_cast<Instruction>(LOr))
+          setExplicitlyUnknownBranchWeightsIfProfiled(*I, DEBUG_TYPE);
       return replaceInstUsesWith(SI, LOr);
     }
 
@@ -3938,10 +3944,16 @@ Instruction *InstCombinerImpl::foldSelectOfBools(SelectInst &SI) {
     if (match(CondVal, m_OneUse(m_Select(m_Value(A), m_Value(B), m_Zero()))) &&
         impliesPoisonOrCond(TrueVal, B, /*Expected=*/true, SQ)) {
       // (A && B) && C --> A && (B & C)
-      Value *LAnd = Builder.CreateLogicalAnd(A, Builder.CreateAnd(B, TrueVal));
-      if (auto *I = dyn_cast<Instruction>(LAnd)) {
-        setExplicitlyUnknownBranchWeightsIfProfiled(*I, DEBUG_TYPE);
-      }
+      // CondVal is select A, B, false, so its weights describe A.
+      Instruction *MDFrom = nullptr;
+      if (!ProfcheckDisableMetadataFixes &&
+          cast<Instruction>(CondVal)->getMetadata(LLVMContext::MD_prof))
+        MDFrom = cast<Instruction>(CondVal);
+      Value *LAnd = Builder.CreateLogicalAnd(A, Builder.CreateAnd(B, TrueVal),
+                                             "", MDFrom);
+      if (!ProfcheckDisableMetadataFixes && !MDFrom)
+        if (auto *I = dyn_cast<Instruction>(LAnd))
+          setExplicitlyUnknownBranchWeightsIfProfiled(*I, DEBUG_TYPE);
       return replaceInstUsesWith(SI, LAnd);
     }
 
diff --git a/llvm/test/Transforms/InstCombine/select-logical-reassoc.ll b/llvm/test/Transforms/InstCombine/select-logical-reassoc.ll
new file mode 100644
index 0000000000000..0023147a32d13
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/select-logical-reassoc.ll
@@ -0,0 +1,56 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+
+define i1 @test_reassoc_or(i1 %a, i1 %b) {
+; CHECK-LABEL: define i1 @test_reassoc_or(
+; CHECK-SAME: i1 [[A:%.*]], i1 [[B:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = select i1 [[A]], i1 true, i1 [[B]], !prof [[PROF0:![0-9]+]]
+; CHECK-NEXT:    ret i1 [[RESULT]]
+;
+  %inner = select i1 %a, i1 true, i1 %b, !prof !0
+  %result = select i1 %inner, i1 true, i1 %b, !prof !1
+  ret i1 %result
+}
+
+define i1 @test_reassoc_and(i1 %a, i1 %b) {
+; CHECK-LABEL: define i1 @test_reassoc_and(
+; CHECK-SAME: i1 [[A:%.*]], i1 [[B:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = select i1 [[A]], i1 [[B]], i1 false, !prof [[PROF0]]
+; CHECK-NEXT:    ret i1 [[RESULT]]
+;
+  %inner = select i1 %a, i1 %b, i1 false, !prof !0
+  %result = select i1 %inner, i1 %b, i1 false, !prof !1
+  ret i1 %result
+}
+
+; Neither select has a profile, but the function does, so the result is unknown.
+define i1 @test_reassoc_or_no_profile(i1 %a, i1 %b) !prof !2 {
+; CHECK-LABEL: define i1 @test_reassoc_or_no_profile(
+; CHECK-SAME: i1 [[A:%.*]], i1 [[B:%.*]]) !prof [[PROF1:![0-9]+]] {
+; CHECK-NEXT:    [[RESULT:%.*]] = select i1 [[A]], i1 true, i1 [[B]], !prof [[PROF2:![0-9]+]]
+; CHECK-NEXT:    ret i1 [[RESULT]]
+;
+  %inner = select i1 %a, i1 true, i1 %b
+  %result = select i1 %inner, i1 true, i1 %b
+  ret i1 %result
+}
+
+define i1 @test_reassoc_and_no_profile(i1 %a, i1 %b) !prof !2 {
+; CHECK-LABEL: define i1 @test_reassoc_and_no_profile(
+; CHECK-SAME: i1 [[A:%.*]], i1 [[B:%.*]]) !prof [[PROF1]] {
+; CHECK-NEXT:    [[RESULT:%.*]] = select i1 [[A]], i1 [[B]], i1 false, !prof [[PROF2]]
+; CHECK-NEXT:    ret i1 [[RESULT]]
+;
+  %inner = select i1 %a, i1 %b, i1 false
+  %result = select i1 %inner, i1 %b, i1 false
+  ret i1 %result
+}
+
+!0 = !{!"branch_weights", i32 3, i32 7}
+!1 = !{!"branch_weights", i32 9, i32 1}
+!2 = !{!"function_entry_count", i64 100}
+;.
+; CHECK: [[PROF0]] = !{!"branch_weights", i32 3, i32 7}
+; CHECK: [[PROF1]] = !{!"function_entry_count", i64 100}
+; CHECK: [[PROF2]] = !{!"unknown", !"instcombine"}
+;.

[nikic]

Why this check? Doesn't this just break !unpredictable copying?