diff --git a/tools/klockstat.py b/tools/klockstat.py
index e2047880370a..540dd4e7d1b9 100755
--- a/tools/klockstat.py
+++ b/tools/klockstat.py
@@ -124,7 +124,7 @@ def stack_id_err(stack_id):
     return 1;
 }
 
-int mutex_lock_enter(struct pt_regs *ctx)
+static int do_mutex_lock_enter(void *ctx, int skip)
 {
     if (!is_enabled())
         return 0;
@@ -149,7 +149,7 @@ def stack_id_err(stack_id):
             return 0;
     }
 
-    int stackid = stack_traces.get_stackid(ctx, 0);
+    int stackid = stack_traces.get_stackid(ctx, skip);
     struct depth_id did = {
       .id    = id,
       .depth = *depth,
@@ -237,7 +237,7 @@ def stack_id_err(stack_id):
     }
 }
 
-int mutex_lock_return(struct pt_regs *ctx)
+static int do_mutex_lock_return(void)
 {
     if (!is_enabled())
         return 0;
@@ -285,7 +285,7 @@ def stack_id_err(stack_id):
     return 0;
 }
 
-int mutex_unlock_enter(struct pt_regs *ctx)
+static int do_mutex_unlock_enter(void)
 {
     if (!is_enabled())
         return 0;
@@ -330,9 +330,49 @@ def stack_id_err(stack_id):
     time_held.delete(&did);
     return 0;
 }
+"""
+
+program_kprobe = """
+int mutex_unlock_enter(struct pt_regs *ctx)
+{
+    return do_mutex_unlock_enter();
+}
+
+int mutex_lock_return(struct pt_regs *ctx)
+{
+    return do_mutex_lock_return();
+}
+
+int mutex_lock_enter(struct pt_regs *ctx)
+{
+    return do_mutex_lock_enter(ctx, 0);
+}
+"""
+
+program_kfunc = """
+KFUNC_PROBE(mutex_unlock, void *lock)
+{
+    do_mutex_unlock_enter();
+}
+
+KRETFUNC_PROBE(mutex_lock, void *lock, int ret)
+{
+    do_mutex_lock_return();
+}
+
+KFUNC_PROBE(mutex_lock, void *lock)
+{
+    do_mutex_lock_enter(ctx, 3);
+}
 
 """
 
+is_support_kfunc = BPF.support_kfunc()
+if is_support_kfunc:
+    program += program_kfunc
+else:
+    program += program_kprobe
+
 def sort_list(maxs, totals, counts):
     if (not args.sort):
         return maxs;
@@ -387,9 +427,10 @@ def display(sort, maxs, totals, counts):
 
 b = BPF(text=program)
 
-b.attach_kprobe(event="mutex_unlock", fn_name="mutex_unlock_enter")
-b.attach_kretprobe(event="mutex_lock", fn_name="mutex_lock_return")
-b.attach_kprobe(event="mutex_lock", fn_name="mutex_lock_enter")
+if not is_support_kfunc:
+    b.attach_kprobe(event="mutex_unlock",  fn_name="mutex_unlock_enter")
+    b.attach_kretprobe(event="mutex_lock", fn_name="mutex_lock_return")
+    b.attach_kprobe(event="mutex_lock",    fn_name="mutex_lock_enter")
 
 enabled = b.get_table("enabled");