[licm] don't drop MD_prof when dropping other metadata (#152420)

Part of Issue #147390

Author: mtrofin

llvm/include/llvm/IR/Instruction.h

@@ -584,9 +584,10 @@ class Instruction : public User,
   dropUBImplyingAttrsAndUnknownMetadata(ArrayRef<unsigned> KnownIDs = {});
 
   /// Drop any attributes or metadata that can cause immediate undefined
-  /// behavior. Retain other attributes/metadata on a best-effort basis.
-  /// This should be used when speculating instructions.
-  LLVM_ABI void dropUBImplyingAttrsAndMetadata();
+  /// behavior. Retain other attributes/metadata on a best-effort basis, as well
+  /// as those passed in `Keep`. This should be used when speculating
+  /// instructions.
+  LLVM_ABI void dropUBImplyingAttrsAndMetadata(ArrayRef<unsigned> Keep = {});
 
   /// Return true if this instruction has UB-implying attributes
   /// that can cause immediate undefined behavior.

llvm/lib/IR/Instruction.cpp

@@ -552,14 +552,19 @@ void Instruction::dropUBImplyingAttrsAndUnknownMetadata(
   CB->removeRetAttrs(UBImplyingAttributes);
 }
 
-void Instruction::dropUBImplyingAttrsAndMetadata() {
+void Instruction::dropUBImplyingAttrsAndMetadata(ArrayRef<unsigned> Keep) {
   // !annotation metadata does not impact semantics.
   // !range, !nonnull and !align produce poison, so they are safe to speculate.
   // !noundef and various AA metadata must be dropped, as it generally produces
   // immediate undefined behavior.
-  unsigned KnownIDs[] = {LLVMContext::MD_annotation, LLVMContext::MD_range,
-                         LLVMContext::MD_nonnull, LLVMContext::MD_align};
-  dropUBImplyingAttrsAndUnknownMetadata(KnownIDs);
+  static const unsigned KnownIDs[] = {
+      LLVMContext::MD_annotation, LLVMContext::MD_range,
+      LLVMContext::MD_nonnull, LLVMContext::MD_align};
+  SmallVector<unsigned> KeepIDs;
+  KeepIDs.reserve(Keep.size() + std::size(KnownIDs));
+  append_range(KeepIDs, KnownIDs);
+  append_range(KeepIDs, Keep);
+  dropUBImplyingAttrsAndUnknownMetadata(KeepIDs);
 }
 
 bool Instruction::hasUBImplyingAttrs() const {

llvm/lib/Transforms/Scalar/LICM.cpp

@@ -1699,8 +1699,12 @@ static void hoist(Instruction &I, const DominatorTree *DT, const Loop *CurLoop,
       // The check on hasMetadataOtherThanDebugLoc is to prevent us from burning
       // time in isGuaranteedToExecute if we don't actually have anything to
       // drop.  It is a compile time optimization, not required for correctness.
-      !SafetyInfo->isGuaranteedToExecute(I, DT, CurLoop))
-    I.dropUBImplyingAttrsAndMetadata();
+      !SafetyInfo->isGuaranteedToExecute(I, DT, CurLoop)) {
+    if (ProfcheckDisableMetadataFixes)
+      I.dropUBImplyingAttrsAndMetadata();
+    else
+      I.dropUBImplyingAttrsAndMetadata({LLVMContext::MD_prof});
+  }
 
   if (isa<PHINode>(I))
     // Move the new node to the end of the phi list in the destination block.

llvm/test/Transforms/LICM/hoist-profdata.ll

@@ -0,0 +1,45 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals --version 2
+; Test that hoisting conditional branches copies the debug and profiling info
+; metadata from the branch being hoisted.
+; RUN: opt -S -passes=licm %s -o - | FileCheck %s
+
+declare i32 @foo()
+
+; to_hoist should get hoisted, and that should not result
+; in a loss of profiling info
+define i32 @hoist_select(i1 %cond, i32 %a, i32 %b) nounwind {
+; CHECK-LABEL: define i32 @hoist_select
+; CHECK-SAME: (i1 [[COND:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TO_HOIST:%.*]] = select i1 [[COND]], i32 [[A]], i32 [[B]], !prof [[PROF0:![0-9]+]]
+; CHECK-NEXT:    br label [[L0:%.*]]
+; CHECK:       L0:
+; CHECK-NEXT:    [[G:%.*]] = call i32 @foo()
+; CHECK-NEXT:    [[SUM:%.*]] = add i32 [[G]], [[TO_HOIST]]
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[SUM]], 0
+; CHECK-NEXT:    br i1 [[C]], label [[L0]], label [[EXIT:%.*]], !prof [[PROF1:![0-9]+]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[SUM_LCSSA:%.*]] = phi i32 [ [[SUM]], [[L0]] ]
+; CHECK-NEXT:    ret i32 [[SUM_LCSSA]]
+;
+entry:
+  br label %L0
+L0:
+  %g = call i32 @foo()
+  %to_hoist = select i1 %cond, i32 %a, i32 %b, !prof !0
+  %sum = add i32 %g, %to_hoist
+  %c = icmp eq i32 %sum, 0
+  br i1 %c, label %L0, label %exit, !prof !1
+
+exit:
+  ret i32 %sum
+}
+
+!0 = !{!"branch_weights", i32 2, i32 5}
+!1 = !{!"branch_weights", i32 101, i32 189}
+;.
+; CHECK: attributes #[[ATTR0]] = { nounwind }
+;.
+; CHECK: [[PROF0]] = !{!"branch_weights", i32 2, i32 5}
+; CHECK: [[PROF1]] = !{!"branch_weights", i32 101, i32 189}
+;.