[licm] don't drop MD_prof when dropping other metadata

[mtrofin]

Part of Issue #147390

[mtrofin]

• [nfc][pgo] Expose traversal of VP metadata elements #153237
• [licm] don't drop MD_prof when dropping other metadata #152420 👈 (View in Graphite)
• [licm] clone MD_prof when hoisting conditional branch #152232
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Mircea Trofin (mtrofin)

Changes

---

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

4 Files Affected:

• (modified) llvm/include/llvm/IR/Instruction.h (+4-3)
• (modified) llvm/lib/IR/Instruction.cpp (+5-3)
• (modified) llvm/lib/Transforms/Scalar/LICM.cpp (+6-2)
• (added) llvm/test/Transforms/LICM/hoist-profdata.ll (+45)

diff --git a/llvm/include/llvm/IR/Instruction.h b/llvm/include/llvm/IR/Instruction.h
index 5d25804a684ac..2eb4fd36c5b7d 100644
--- a/llvm/include/llvm/IR/Instruction.h
+++ b/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.
diff --git a/llvm/lib/IR/Instruction.cpp b/llvm/lib/IR/Instruction.cpp
index cf19b97333514..2815ea8498b4a 100644
--- a/llvm/lib/IR/Instruction.cpp
+++ b/llvm/lib/IR/Instruction.cpp
@@ -554,13 +554,15 @@ 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};
+  SmallVector<unsigned> KnownIDs(Keep.size() + 4);
+  KnownIDs.append({LLVMContext::MD_annotation, LLVMContext::MD_range,
+                   LLVMContext::MD_nonnull, LLVMContext::MD_align});
+  append_range(KnownIDs, Keep);
   dropUBImplyingAttrsAndUnknownMetadata(KnownIDs);
 }
 
diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
index 34adc7745f67c..7a203638a1d78 100644
--- a/llvm/lib/Transforms/Scalar/LICM.cpp
+++ b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -1698,8 +1698,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.
diff --git a/llvm/test/Transforms/LICM/hoist-profdata.ll b/llvm/test/Transforms/LICM/hoist-profdata.ll
new file mode 100644
index 0000000000000..18fa1b9f92e8a
--- /dev/null
+++ b/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}
+;.

[llvmbot]

@llvm/pr-subscribers-llvm-ir

Author: Mircea Trofin (mtrofin)

Changes

---

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

4 Files Affected:

• (modified) llvm/include/llvm/IR/Instruction.h (+4-3)
• (modified) llvm/lib/IR/Instruction.cpp (+5-3)
• (modified) llvm/lib/Transforms/Scalar/LICM.cpp (+6-2)
• (added) llvm/test/Transforms/LICM/hoist-profdata.ll (+45)

diff --git a/llvm/include/llvm/IR/Instruction.h b/llvm/include/llvm/IR/Instruction.h
index 5d25804a684ac..2eb4fd36c5b7d 100644
--- a/llvm/include/llvm/IR/Instruction.h
+++ b/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.
diff --git a/llvm/lib/IR/Instruction.cpp b/llvm/lib/IR/Instruction.cpp
index cf19b97333514..2815ea8498b4a 100644
--- a/llvm/lib/IR/Instruction.cpp
+++ b/llvm/lib/IR/Instruction.cpp
@@ -554,13 +554,15 @@ 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};
+  SmallVector<unsigned> KnownIDs(Keep.size() + 4);
+  KnownIDs.append({LLVMContext::MD_annotation, LLVMContext::MD_range,
+                   LLVMContext::MD_nonnull, LLVMContext::MD_align});
+  append_range(KnownIDs, Keep);
   dropUBImplyingAttrsAndUnknownMetadata(KnownIDs);
 }
 
diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
index 34adc7745f67c..7a203638a1d78 100644
--- a/llvm/lib/Transforms/Scalar/LICM.cpp
+++ b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -1698,8 +1698,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.
diff --git a/llvm/test/Transforms/LICM/hoist-profdata.ll b/llvm/test/Transforms/LICM/hoist-profdata.ll
new file mode 100644
index 0000000000000..18fa1b9f92e8a
--- /dev/null
+++ b/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}
+;.

[snehasish]

Why do you need to pass Keep.size() + std::size(KnownIDs)? The append_range calls below insert at the end so I think this means you'll end up with a bunch of zeros in the beginning.

[mtrofin]

Because I meant reserve

[snehasish]

Could we directly call dropUBImplyingAttrsAndUnknownMetadata here with the md_prof arg instead of changing the other function with a keep argument?

[mtrofin]

but we'd still want to pass the other bunch of OK metadata, which would mean exposing it.

[snehasish]

Ah. Didn't realize it wasn't exposed already.

[snehasish]

Lgtm

[snehasish]

Is this useful in practice? SmallVector might already allocate enough inline storage to not require an allocation for the small number of metadata types we have.

[mtrofin]

maybe. it's harmless though and at best a tiny efficiency.

[mtrofin]

Merge activity

• Aug 16, 2:24 PM UTC: A user started a stack merge that includes this pull request via Graphite.
• Aug 16, 2:26 PM UTC: @mtrofin merged this pull request with Graphite.

[llvm-ci]

LLVM Buildbot has detected a new failure on builder llvm-clang-aarch64-darwin running on doug-worker-4 while building llvm at step 6 "test-build-unified-tree-check-all".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/190/builds/25508

Here is the relevant piece of the build log for the reference

Step 6 (test-build-unified-tree-check-all) failure: test (failure)
******************** TEST 'LLVM :: ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll' FAILED ********************
Exit Code: 2

Command Output (stderr):
--
/Volumes/RAMDisk/buildbot-root/aarch64-darwin/build/bin/lli -jit-kind=orc-lazy -compile-threads=2 -thread-entry hello /Users/buildbot/buildbot-root/llvm-project/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll | /Volumes/RAMDisk/buildbot-root/aarch64-darwin/build/bin/FileCheck /Users/buildbot/buildbot-root/llvm-project/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll # RUN: at line 1
+ /Volumes/RAMDisk/buildbot-root/aarch64-darwin/build/bin/lli -jit-kind=orc-lazy -compile-threads=2 -thread-entry hello /Users/buildbot/buildbot-root/llvm-project/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll
+ /Volumes/RAMDisk/buildbot-root/aarch64-darwin/build/bin/FileCheck /Users/buildbot/buildbot-root/llvm-project/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
 #0 0x0000000103b08748 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/Volumes/RAMDisk/buildbot-root/aarch64-darwin/build/bin/lli+0x100f0c748)
 #1 0x0000000103b064f8 llvm::sys::RunSignalHandlers() (/Volumes/RAMDisk/buildbot-root/aarch64-darwin/build/bin/lli+0x100f0a4f8)
 #2 0x0000000103b09248 SignalHandler(int, __siginfo*, void*) (/Volumes/RAMDisk/buildbot-root/aarch64-darwin/build/bin/lli+0x100f0d248)
 #3 0x0000000188543584 (/usr/lib/system/libsystem_platform.dylib+0x18047b584)
 #4 0x000000018851221c (/usr/lib/system/libsystem_pthread.dylib+0x18044a21c)
 #5 0x0000000188438ad0 (/usr/lib/libc++.1.dylib+0x180370ad0)
 #6 0x000000010368da4c void llvm::detail::UniqueFunctionBase<void, llvm::Expected<llvm::DenseMap<llvm::orc::SymbolStringPtr, llvm::orc::ExecutorSymbolDef, llvm::DenseMapInfo<llvm::orc::SymbolStringPtr, void>, llvm::detail::DenseMapPair<llvm::orc::SymbolStringPtr, llvm::orc::ExecutorSymbolDef>>>>::CallImpl<llvm::orc::Platform::lookupInitSymbols(llvm::orc::ExecutionSession&, llvm::DenseMap<llvm::orc::JITDylib*, llvm::orc::SymbolLookupSet, llvm::DenseMapInfo<llvm::orc::JITDylib*, void>, llvm::detail::DenseMapPair<llvm::orc::JITDylib*, llvm::orc::SymbolLookupSet>> const&)::$_45>(void*, llvm::Expected<llvm::DenseMap<llvm::orc::SymbolStringPtr, llvm::orc::ExecutorSymbolDef, llvm::DenseMapInfo<llvm::orc::SymbolStringPtr, void>, llvm::detail::DenseMapPair<llvm::orc::SymbolStringPtr, llvm::orc::ExecutorSymbolDef>>>&) (/Volumes/RAMDisk/buildbot-root/aarch64-darwin/build/bin/lli+0x100a91a4c)
 #7 0x000000010368964c llvm::orc::AsynchronousSymbolQuery::handleComplete(llvm::orc::ExecutionSession&)::RunQueryCompleteTask::run() (/Volumes/RAMDisk/buildbot-root/aarch64-darwin/build/bin/lli+0x100a8d64c)
 #8 0x00000001037515e8 void* std::__1::__thread_proxy[abi:un170006]<std::__1::tuple<std::__1::unique_ptr<std::__1::__thread_struct, std::__1::default_delete<std::__1::__thread_struct>>, llvm::orc::DynamicThreadPoolTaskDispatcher::dispatch(std::__1::unique_ptr<llvm::orc::Task, std::__1::default_delete<llvm::orc::Task>>)::$_0>>(void*) (/Volumes/RAMDisk/buildbot-root/aarch64-darwin/build/bin/lli+0x100b555e8)
 #9 0x0000000188512f94 (/usr/lib/system/libsystem_pthread.dylib+0x18044af94)
#10 0x000000018850dd34 (/usr/lib/system/libsystem_pthread.dylib+0x180445d34)
FileCheck error: '<stdin>' is empty.
FileCheck command line:  /Volumes/RAMDisk/buildbot-root/aarch64-darwin/build/bin/FileCheck /Users/buildbot/buildbot-root/llvm-project/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll

--

********************