diff --git a/src/runtime/Events/configs/DepModule_TrackContext_Events.yaml b/src/runtime/Events/configs/DepModule_TrackContext_Events.yaml
index 15ec222..58a146e 100644
--- a/src/runtime/Events/configs/DepModule_TrackContext_Events.yaml
+++ b/src/runtime/Events/configs/DepModule_TrackContext_Events.yaml
@@ -10,6 +10,9 @@ events:
   target_loop_invoc: []
   target_loop_iter: []
   target_loop_exit: []
-  func_entry: [function_id] # optional if tracking context
-  func_exit: [function_id] # optional if tracking context
+  func_call_push: [inst_id] # optional if tracking context
+  func_call_pop: [] # optional if tracking context
   finished: []
+
+  # func_entry: [function_id] # optional if tracking context
+  # func_exit: [function_id] # optional if tracking context
diff --git a/src/runtime/Events/configs/api.yaml b/src/runtime/Events/configs/api.yaml
index 416b2f0..81fd5f6 100644
--- a/src/runtime/Events/configs/api.yaml
+++ b/src/runtime/Events/configs/api.yaml
@@ -47,6 +47,10 @@ events:
     function_id: 32
   func_exit:
     function_id: 32
+  func_call_push:
+    inst_id: 32
+  func_call_pop:
+    # no arg
   points_to_inst:
     inst_id: 32
     ptr: 64
diff --git a/src/runtime/ProfilingModules/DependenceWithContextModule.cpp b/src/runtime/ProfilingModules/DependenceWithContextModule.cpp
index 573ce7c..e94d011 100644
--- a/src/runtime/ProfilingModules/DependenceWithContextModule.cpp
+++ b/src/runtime/ProfilingModules/DependenceWithContextModule.cpp
@@ -191,13 +191,15 @@ void DependenceWithContextModule::loop_exit() __attribute__((always_inline)) {
 }
 
 void DependenceWithContextModule::func_entry(uint32_t instr) {
-  if (nested_level == 1) {
+  if(callgraph_level == 0) {
     context = instr;
   }
+  callgraph_level++;
 }
 
 void DependenceWithContextModule::func_exit(uint32_t instr) {
-  if (nested_level == 1) {
+  callgraph_level--;
+  if(callgraph_level == 0) {
     context = 0;
   }
 }
diff --git a/src/runtime/ProfilingModules/DependenceWithContextModule.h b/src/runtime/ProfilingModules/DependenceWithContextModule.h
index 9ef85d3..da59e3f 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 callgraph_level = 0;
 
 #ifdef COLLECT_TRACE
   // Collect trace
diff --git a/src/runtime/SLAMPcustom/consumer/consumer.cpp b/src/runtime/SLAMPcustom/consumer/consumer.cpp
index 9d18eee..cd23d3a 100644
--- a/src/runtime/SLAMPcustom/consumer/consumer.cpp
+++ b/src/runtime/SLAMPcustom/consumer/consumer.cpp
@@ -56,6 +56,8 @@ enum class UnifiedAction : char {
   FUNC_EXIT,
   POINTS_TO_INST,
   POINTS_TO_ARG,
+  FUNC_CALL_PUSH,
+  FUNC_CALL_POP,
   FINISHED
 };
 
@@ -1273,7 +1275,8 @@ void consume_loop(DoubleQueue &dq,
 }
 
 void consume_loop_dep_with_context(DoubleQueue &dq,
-                  DependenceWithContextModule &depMod) CONSUME_LOOP_ATTRIBUTES {
+                                   DependenceWithContextModule &depMod)
+    CONSUME_LOOP_ATTRIBUTES {
   uint64_t rdtsc_start = 0;
   uint64_t counter = 0;
   uint32_t loop_id;
@@ -1414,6 +1417,17 @@ void consume_loop_dep_with_context(DoubleQueue &dq,
       depMod.func_exit(func_id);
       break;
     };
+    case Action::FUNC_CALL_PUSH: {
+      uint32_t caller_id;
+      dq.unpack_32(caller_id);
+      depMod.func_entry(caller_id);
+      break;
+    };
+    case Action::FUNC_CALL_POP: {
+      depMod.func_exit(0);
+      break;
+    };
+
 #ifdef UNIFIED_WORKFLOW
     case Action::FREE:
     case Action::LOOP_ENTRY:
@@ -1718,7 +1732,10 @@ int main(int argc, char **argv) {
       std::cout << "Running in " << THREAD_COUNT << " threads" << std::endl;
       for (unsigned i = 0; i < THREAD_COUNT; i++) {
         threads.emplace_back(
-            [&](unsigned id) { consume_loop_dep_with_context(*dqs[id], *depMods[id]); }, i);
+            [&](unsigned id) {
+              consume_loop_dep_with_context(*dqs[id], *depMods[id]);
+            },
+            i);
       }
 
       for (auto &t : threads) {
diff --git a/src/runtime/frontend/custom_produce.h b/src/runtime/frontend/custom_produce.h
index 70da152..fb4ab90 100644
--- a/src/runtime/frontend/custom_produce.h
+++ b/src/runtime/frontend/custom_produce.h
@@ -23,6 +23,8 @@ enum UnifiedAction : char {
   FUNC_EXIT,
   POINTS_TO_INST,
   POINTS_TO_ARG,
+  FUNC_CALL_PUSH,
+  FUNC_CALL_POP,
   FINISHED
 };
 
diff --git a/src/runtime/frontend/frontend.cpp b/src/runtime/frontend/frontend.cpp
index 9bde88b..ad77a9f 100644
--- a/src/runtime/frontend/frontend.cpp
+++ b/src/runtime/frontend/frontend.cpp
@@ -102,6 +102,14 @@ static uint32_t ext_fn_inst_id = 0;
 #define PRODUCE_STORE(size, instr, addr)
 #endif
 
+#ifndef PRODUCE_FUNC_CALL_PUSH
+#define PRODUCE_FUNC_CALL_PUSH(inst_id)
+#endif
+
+#ifndef PRODUCE_FUNC_CALL_POP
+#define PRODUCE_FUNC_CALL_POP()
+#endif
+
 static volatile char *lc_dummy = NULL;
 
 PRODUCE_QUEUE_DEFINE();
@@ -342,8 +350,9 @@ void memalign_callback(void *ptr, size_t alignment, size_t size) {
 
 // FIXME: a bunch of unused functions
 void SLAMP_main_entry(uint32_t argc, char **argv, char **env) {}
-void SLAMP_push(const uint32_t instr) {}
-void SLAMP_pop() {}
+void SLAMP_push(const uint32_t instr) { if(on_profiling) PRODUCE_FUNC_CALL_PUSH(instr); }
+void SLAMP_pop() { if(on_profiling) PRODUCE_FUNC_CALL_POP(); }
+
 void SLAMP_allocated(uint64_t addr) {}
 
 void SLAMP_callback_stack_alloca(uint32_t instr, void *ptr, uint64_t array_size,