[llvm] 6827c4f - [MemProf] Add helper to access the back (last) call stack id

[Teresa Johnson via llvm-commits]

Author: Teresa Johnson
Date: 2023-02-03T07:51:32-08:00
New Revision: 6827c4f0dea1e3840ace464a11328199e11fb163

URL: https://github.com/llvm/llvm-project/commit/6827c4f0dea1e3840ace464a11328199e11fb163
DIFF: https://github.com/llvm/llvm-project/commit/6827c4f0dea1e3840ace464a11328199e11fb163.diff

LOG: [MemProf] Add helper to access the back (last) call stack id

This is split out of D140908 as suggested.

Differential Revision: https://reviews.llvm.org/D143184

Added: 
    

Modified: 
    llvm/include/llvm/Analysis/MemoryProfileInfo.h
    llvm/lib/Analysis/MemoryProfileInfo.cpp
    llvm/unittests/Analysis/MemoryProfileInfoTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Analysis/MemoryProfileInfo.h b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
index 24956e7815722..2ca396bf1453f 100644
--- a/llvm/include/llvm/Analysis/MemoryProfileInfo.h
+++ b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
@@ -128,6 +128,7 @@ template <class NodeT, class IteratorT> class CallStack {
   CallStackIterator begin() const;
   CallStackIterator end() const { return CallStackIterator(N, /*End*/ true); }
   CallStackIterator beginAfterSharedPrefix(CallStack &Other);
+  uint64_t back() const;
 
 private:
   const NodeT *N = nullptr;
@@ -137,8 +138,10 @@ template <class NodeT, class IteratorT>
 CallStack<NodeT, IteratorT>::CallStackIterator::CallStackIterator(
     const NodeT *N, bool End)
     : N(N) {
-  if (!N)
+  if (!N) {
+    Iter = nullptr;
     return;
+  }
   Iter = End ? N->StackIdIndices.end() : N->StackIdIndices.begin();
 }
 
@@ -148,6 +151,12 @@ uint64_t CallStack<NodeT, IteratorT>::CallStackIterator::operator*() {
   return *Iter;
 }
 
+template <class NodeT, class IteratorT>
+uint64_t CallStack<NodeT, IteratorT>::back() const {
+  assert(N);
+  return N->StackIdIndices.back();
+}
+
 template <class NodeT, class IteratorT>
 typename CallStack<NodeT, IteratorT>::CallStackIterator
 CallStack<NodeT, IteratorT>::begin() const {
@@ -170,6 +179,7 @@ CallStack<MDNode, MDNode::op_iterator>::CallStackIterator::CallStackIterator(
     const MDNode *N, bool End);
 template <>
 uint64_t CallStack<MDNode, MDNode::op_iterator>::CallStackIterator::operator*();
+template <> uint64_t CallStack<MDNode, MDNode::op_iterator>::back() const;
 
 } // end namespace memprof
 } // end namespace llvm

diff  --git a/llvm/lib/Analysis/MemoryProfileInfo.cpp b/llvm/lib/Analysis/MemoryProfileInfo.cpp
index 8ced1d2fd140c..2f36ece35195c 100644
--- a/llvm/lib/Analysis/MemoryProfileInfo.cpp
+++ b/llvm/lib/Analysis/MemoryProfileInfo.cpp
@@ -242,3 +242,9 @@ CallStack<MDNode, MDNode::op_iterator>::CallStackIterator::operator*() {
   assert(StackIdCInt);
   return StackIdCInt->getZExtValue();
 }
+
+template <> uint64_t CallStack<MDNode, MDNode::op_iterator>::back() const {
+  assert(N);
+  return mdconst::dyn_extract<ConstantInt>(N->operands().back())
+      ->getZExtValue();
+}

diff  --git a/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp b/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp
index ffe822ebc5b46..096295eb3f2d6 100644
--- a/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp
+++ b/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp
@@ -401,6 +401,8 @@ declare noundef nonnull ptr @_Znam(i64 noundef)
     auto *MIBMD = cast<const MDNode>(MIBOp);
     MDNode *StackNode = getMIBStackNode(MIBMD);
     CallStack<MDNode, MDNode::op_iterator> StackContext(StackNode);
+    uint64_t ExpectedBack = First ? 4 : 5;
+    EXPECT_EQ(StackContext.back(), ExpectedBack);
     std::vector<uint64_t> StackIds;
     for (auto ContextIter = StackContext.beginAfterSharedPrefix(InstCallsite);
          ContextIter != StackContext.end(); ++ContextIter)
@@ -450,6 +452,8 @@ TEST_F(MemoryProfileInfoTest, CallStackTestSummary) {
     for (auto &MIB : AI.MIBs) {
       CallStack<MIBInfo, SmallVector<unsigned>::const_iterator> StackContext(
           &MIB);
+      uint64_t ExpectedBack = First ? 4 : 5;
+      EXPECT_EQ(Index->getStackIdAtIndex(StackContext.back()), ExpectedBack);
       std::vector<uint64_t> StackIds;
       for (auto StackIdIndex : StackContext)
         StackIds.push_back(Index->getStackIdAtIndex(StackIdIndex));