From cd0064bac94bdf11bc917f0113b5f6d9b7d368db Mon Sep 17 00:00:00 2001
From: Yebin Chon <yebin.chon@princeton.edu>
Date: Mon, 8 Apr 2024 21:12:14 -0400
Subject: [PATCH] bugfix to track caller inst

---
 .../DependenceWithContextModule.cpp           | 24 ++++++++++++++-----
 .../DependenceWithContextModule.h             |  8 +++++--
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/src/runtime/ProfilingModules/DependenceWithContextModule.cpp b/src/runtime/ProfilingModules/DependenceWithContextModule.cpp
index 573ce7c..158ffdf 100644
--- a/src/runtime/ProfilingModules/DependenceWithContextModule.cpp
+++ b/src/runtime/ProfilingModules/DependenceWithContextModule.cpp
@@ -190,16 +190,28 @@ void DependenceWithContextModule::loop_exit() __attribute__((always_inline)) {
   nested_level--;
 }
 
-void DependenceWithContextModule::func_entry(uint32_t instr) {
-  if (nested_level == 1) {
+//void DependenceWithContextModule::func_entry(uint32_t instr) {
+//  if (nested_level == 1) {
+//    context = instr;
+//  }
+//}
+//
+//void DependenceWithContextModule::func_exit(uint32_t instr) {
+//  if (nested_level == 1) {
+//    context = 0;
+//  }
+//}
+
+void DependenceWithContextModule::push(uint32_t instr) __attribute__((always_inline)) {
+  if(func_level == 0)
     context = instr;
-  }
+  func_level++;
 }
 
-void DependenceWithContextModule::func_exit(uint32_t instr) {
-  if (nested_level == 1) {
+void DependenceWithContextModule::pop() __attribute__((always_inline)) {
+  func_level--;
+  if(func_level == 0)
     context = 0;
-  }
 }
 
 void DependenceWithContextModule::merge_dep(DependenceWithContextModule &other) {
diff --git a/src/runtime/ProfilingModules/DependenceWithContextModule.h b/src/runtime/ProfilingModules/DependenceWithContextModule.h
index 9ef85d3..c07a5ff 100644
--- a/src/runtime/ProfilingModules/DependenceWithContextModule.h
+++ b/src/runtime/ProfilingModules/DependenceWithContextModule.h
@@ -37,6 +37,7 @@ class DependenceWithContextModule : public LocalWriteModule {
 
   unsigned int context = 0;
   int nested_level = 0;
+  int func_level = 0;
 
 #ifdef COLLECT_TRACE
   // Collect trace
@@ -80,8 +81,11 @@ class DependenceWithContextModule : public LocalWriteModule {
   void loop_invoc();
   void loop_iter();
   void loop_exit();
-  void func_entry(uint32_t context);
-  void func_exit(uint32_t context);
+  //void func_entry(uint32_t context);
+  //void func_exit(uint32_t context);
+  //FIXME: do we also need ext_push/pop?
+  void push(uint32_t instr);
+  void pop();
 
   void merge_dep(DependenceWithContextModule &other);
 };