Merge branch 'main' into main

Author: brandonxin

llvm/include/llvm/ADT/DenseMap.h

@@ -1184,10 +1184,14 @@ class DenseMapIterator : DebugEpochBase::HandleBase {
   using iterator_category = std::forward_iterator_tag;
 
 private:
-  pointer Ptr = nullptr;
-  pointer End = nullptr;
+  using BucketItTy =
+      std::conditional_t<shouldReverseIterate<KeyT>(),
+                         std::reverse_iterator<pointer>, pointer>;
 
-  DenseMapIterator(pointer Pos, pointer E, const DebugEpochBase &Epoch)
+  BucketItTy Ptr = {};
+  BucketItTy End = {};
+
+  DenseMapIterator(BucketItTy Pos, BucketItTy E, const DebugEpochBase &Epoch)
       : DebugEpochBase::HandleBase(&Epoch), Ptr(Pos), End(E) {
     assert(isHandleInSync() && "invalid construction!");
   }
@@ -1201,29 +1205,24 @@ class DenseMapIterator : DebugEpochBase::HandleBase {
     // empty buckets.
     if (IsEmpty)
       return makeEnd(Buckets, Epoch);
-    if (shouldReverseIterate<KeyT>()) {
-      DenseMapIterator Iter(Buckets.end(), Buckets.begin(), Epoch);
-      Iter.RetreatPastEmptyBuckets();
-      return Iter;
-    }
-    DenseMapIterator Iter(Buckets.begin(), Buckets.end(), Epoch);
+    auto R = maybeReverse(Buckets);
+    DenseMapIterator Iter(R.begin(), R.end(), Epoch);
     Iter.AdvancePastEmptyBuckets();
     return Iter;
   }
 
   static DenseMapIterator makeEnd(iterator_range<pointer> Buckets,
                                   const DebugEpochBase &Epoch) {
-    if (shouldReverseIterate<KeyT>())
-      return DenseMapIterator(Buckets.begin(), Buckets.begin(), Epoch);
-    return DenseMapIterator(Buckets.end(), Buckets.end(), Epoch);
+    auto R = maybeReverse(Buckets);
+    return DenseMapIterator(R.end(), R.end(), Epoch);
   }
 
   static DenseMapIterator makeIterator(pointer P,
                                        iterator_range<pointer> Buckets,
                                        const DebugEpochBase &Epoch) {
-    if (shouldReverseIterate<KeyT>())
-      return DenseMapIterator(P + 1, Buckets.begin(), Epoch);
-    return DenseMapIterator(P, Buckets.end(), Epoch);
+    auto R = maybeReverse(Buckets);
+    constexpr int Offset = shouldReverseIterate<KeyT>() ? 1 : 0;
+    return DenseMapIterator(BucketItTy(P + Offset), R.end(), Epoch);
   }
 
   // Converting ctor from non-const iterators to const iterators. SFINAE'd out
@@ -1238,16 +1237,16 @@ class DenseMapIterator : DebugEpochBase::HandleBase {
   reference operator*() const {
     assert(isHandleInSync() && "invalid iterator access!");
     assert(Ptr != End && "dereferencing end() iterator");
-    if (shouldReverseIterate<KeyT>())
-      return Ptr[-1];
     return *Ptr;
   }
   pointer operator->() const { return &operator*(); }
 
   friend bool operator==(const DenseMapIterator &LHS,
                          const DenseMapIterator &RHS) {
-    assert((!LHS.Ptr || LHS.isHandleInSync()) && "handle not in sync!");
-    assert((!RHS.Ptr || RHS.isHandleInSync()) && "handle not in sync!");
+    assert((!LHS.getEpochAddress() || LHS.isHandleInSync()) &&
+           "handle not in sync!");
+    assert((!RHS.getEpochAddress() || RHS.isHandleInSync()) &&
+           "handle not in sync!");
     assert(LHS.getEpochAddress() == RHS.getEpochAddress() &&
            "comparing incomparable iterators!");
     return LHS.Ptr == RHS.Ptr;
@@ -1261,11 +1260,6 @@ class DenseMapIterator : DebugEpochBase::HandleBase {
   inline DenseMapIterator &operator++() { // Preincrement
     assert(isHandleInSync() && "invalid iterator access!");
     assert(Ptr != End && "incrementing end() iterator");
-    if (shouldReverseIterate<KeyT>()) {
-      --Ptr;
-      RetreatPastEmptyBuckets();
-      return *this;
-    }
     ++Ptr;
     AdvancePastEmptyBuckets();
     return *this;
@@ -1288,14 +1282,11 @@ class DenseMapIterator : DebugEpochBase::HandleBase {
       ++Ptr;
   }
 
-  void RetreatPastEmptyBuckets() {
-    assert(Ptr >= End);
-    const KeyT Empty = KeyInfoT::getEmptyKey();
-    const KeyT Tombstone = KeyInfoT::getTombstoneKey();
-
-    while (Ptr != End && (KeyInfoT::isEqual(Ptr[-1].getFirst(), Empty) ||
-                          KeyInfoT::isEqual(Ptr[-1].getFirst(), Tombstone)))
-      --Ptr;
+  static auto maybeReverse(iterator_range<pointer> Range) {
+    if constexpr (shouldReverseIterate<KeyT>())
+      return reverse(Range);
+    else
+      return Range;
   }
 };
 

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

@@ -89,6 +89,8 @@ static cl::opt<bool> SpecializeLiteralConstant(
         "Enable specialization of functions that take a literal constant as an "
         "argument"));
 
+extern cl::opt<bool> ProfcheckDisableMetadataFixes;
+
 bool InstCostVisitor::canEliminateSuccessor(BasicBlock *BB,
                                             BasicBlock *Succ) const {
   unsigned I = 0;
@@ -784,9 +786,31 @@ bool FunctionSpecializer::run() {
 
     // Update the known call sites to call the clone.
     for (CallBase *Call : S.CallSites) {
+      Function *Clone = S.Clone;
       LLVM_DEBUG(dbgs() << "FnSpecialization: Redirecting " << *Call
-                        << " to call " << S.Clone->getName() << "\n");
+                        << " to call " << Clone->getName() << "\n");
       Call->setCalledFunction(S.Clone);
+      auto &BFI = GetBFI(*Call->getFunction());
+      std::optional<uint64_t> Count =
+          BFI.getBlockProfileCount(Call->getParent());
+      if (Count && !ProfcheckDisableMetadataFixes) {
+        std::optional<llvm::Function::ProfileCount> MaybeCloneCount =
+            Clone->getEntryCount();
+        assert(MaybeCloneCount && "Clone entry count was not set!");
+        uint64_t CallCount = *Count + MaybeCloneCount->getCount();
+        Clone->setEntryCount(CallCount);
+        if (std::optional<llvm::Function::ProfileCount> MaybeOriginalCount =
+                S.F->getEntryCount()) {
+          uint64_t OriginalCount = MaybeOriginalCount->getCount();
+          if (OriginalCount >= CallCount) {
+            S.F->setEntryCount(OriginalCount - CallCount);
+          } else {
+            // This should generally not happen as that would mean there are
+            // more computed calls to the function than what was recorded.
+            LLVM_DEBUG(S.F->setEntryCount(0));
+          }
+        }
+      }
     }
 
     Clones.push_back(S.Clone);
@@ -1043,6 +1067,9 @@ Function *FunctionSpecializer::createSpecialization(Function *F,
   // clone must.
   Clone->setLinkage(GlobalValue::InternalLinkage);
 
+  if (F->getEntryCount() && !ProfcheckDisableMetadataFixes)
+    Clone->setEntryCount(0);
+
   // Initialize the lattice state of the arguments of the function clone,
   // marking the argument on which we specialized the function constant
   // with the given value.

llvm/test/Transforms/FunctionSpecialization/profile-counts.ll

@@ -0,0 +1,52 @@
+; RUN: opt -passes="ipsccp<func-spec>" -force-specialization -S < %s | FileCheck %s
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+
+@A = external dso_local constant i32, align 4
+@B = external dso_local constant i32, align 4
+
+; CHECK: define dso_local i32 @bar(i32 %x, i32 %y, ptr %z) !prof ![[BAR_PROF:[0-9]]] {
+define dso_local i32 @bar(i32 %x, i32 %y, ptr %z) !prof !0 {
+entry:
+  %tobool = icmp ne i32 %x, 0
+; CHECK: br i1 %tobool, label %if.then, label %if.else, !prof ![[BRANCH_PROF:[0-9]]]
+  br i1 %tobool, label %if.then, label %if.else, !prof !1
+
+; CHECK: if.then:
+; CHECK: call i32 @foo.specialized.1(i32 %x, ptr @A)
+if.then:
+  %call = call i32 @foo(i32 %x, ptr @A)
+  br label %return
+
+; CHECK: if.else:
+; CHECK: call i32 @foo.specialized.2(i32 %y, ptr @B)
+if.else:
+  %call1 = call i32 @foo(i32 %y, ptr @B)
+  br label %return
+
+; CHECK: return:
+; CHECK: %call2 = call i32 @foo(i32 %x, ptr %z)
+return:
+  %retval.0 = phi i32 [ %call, %if.then ], [ %call1, %if.else ]
+  %call2 = call i32 @foo(i32 %x, ptr %z);
+  %add = add i32 %retval.0, %call2
+  ret i32 %add
+}
+
+; CHECK: define internal i32 @foo(i32 %x, ptr %b) !prof ![[FOO_UNSPEC_PROF:[0-9]]]
+; CHECK: define internal i32 @foo.specialized.1(i32 %x, ptr %b) !prof ![[FOO_SPEC_1_PROF:[0-9]]]
+; CHECK: define internal i32 @foo.specialized.2(i32 %x, ptr %b) !prof ![[FOO_SPEC_2_PROF:[0-9]]]
+define internal i32 @foo(i32 %x, ptr %b) !prof !2 {
+entry:
+  %0 = load i32, ptr %b, align 4
+  %add = add nsw i32 %x, %0
+  ret i32 %add
+}
+
+; CHECK: ![[BAR_PROF]] = !{!"function_entry_count", i64 1000}
+; CHECK: ![[BRANCH_PROF]] = !{!"branch_weights", i32 1, i32 3}
+; CHECK: ![[FOO_UNSPEC_PROF]] =  !{!"function_entry_count", i64 234}
+; CHECK: ![[FOO_SPEC_1_PROF]] = !{!"function_entry_count", i64 250}
+; CHECK: ![[FOO_SPEC_2_PROF]] = !{!"function_entry_count", i64 750}
+!0 = !{!"function_entry_count", i64 1000}
+!1 = !{!"branch_weights", i32 1, i32 3}
+!2 = !{!"function_entry_count", i64 1234}

mlir/lib/Conversion/MPIToLLVM/MPIToLLVM.cpp

@@ -405,7 +405,8 @@ std::unique_ptr<MPIImplTraits> MPIImplTraits::get(ModuleOp &moduleOp) {
     return std::make_unique<OMPIImplTraits>(moduleOp);
   if (!strAttr || strAttr.getValue() != "MPICH")
     moduleOp.emitWarning() << "Unknown \"MPI:Implementation\" value in DLTI ("
-                           << strAttr.getValue() << "), defaulting to MPICH";
+                           << (strAttr ? strAttr.getValue() : "<NULL>")
+                           << "), defaulting to MPICH";
   return std::make_unique<MPICHImplTraits>(moduleOp);
 }