diff --git a/include/trace/events/sched.h b/include/trace/events/sched.h
index 028f49662ac34f53e0cb452a330f2ebe0857eaad..eb5ec1fb66b4b365c62ab7a8f8431420b960910f 100644
--- a/include/trace/events/sched.h
+++ b/include/trace/events/sched.h
@@ -705,6 +705,33 @@ DECLARE_TRACE(sched_update_nr_running_tp,
 	TP_PROTO(struct rq *rq, int change),
 	TP_ARGS(rq, change));
 
+DECLARE_EVENT_CLASS(psi_memstall_template,
+
+	TP_PROTO(unsigned long function),
+
+	TP_ARGS(function),
+
+	TP_STRUCT__entry(
+		__field(unsigned long, function)
+	),
+
+	TP_fast_assign(
+		__entry->function = function;
+	),
+
+	TP_printk("%ps", (void *)__entry->function)
+);
+
+DEFINE_EVENT(psi_memstall_template, psi_memstall_enter,
+	TP_PROTO(unsigned long function),
+	TP_ARGS(function)
+);
+
+DEFINE_EVENT(psi_memstall_template, psi_memstall_leave,
+	TP_PROTO(unsigned long function),
+	TP_ARGS(function)
+);
+
 #endif /* _TRACE_SCHED_H */
 
 /* This part must be outside protection */
diff --git a/kernel/sched/psi.c b/kernel/sched/psi.c
index c84c6f7b59fd5e68dde741d7d9c0d89edf1d3d7b..7a5a898ceb6f18706baabbe2813825aa30ff51e5 100644
--- a/kernel/sched/psi.c
+++ b/kernel/sched/psi.c
@@ -142,6 +142,8 @@
 #include <linux/psi.h>
 #include "sched.h"
 
+#include <trace/events/sched.h>
+
 static int psi_bug __read_mostly;
 
 DEFINE_STATIC_KEY_FALSE(psi_disabled);
@@ -904,6 +906,8 @@ void psi_memstall_enter(unsigned long *flags)
 	*flags = current->in_memstall;
 	if (*flags)
 		return;
+
+	trace_psi_memstall_enter(_RET_IP_);
 	/*
 	 * in_memstall setting & accounting needs to be atomic wrt
 	 * changes to the task's scheduling state, otherwise we can
@@ -933,6 +937,8 @@ void psi_memstall_leave(unsigned long *flags)
 
 	if (*flags)
 		return;
+
+	trace_psi_memstall_leave(_RET_IP_);
 	/*
 	 * in_memstall clearing & accounting needs to be atomic wrt
 	 * changes to the task's scheduling state, otherwise we could