diff --git a/kernel/trace/ftrace.c b/kernel/trace/ftrace.c
index ad4440da5b78ce14be4135d11cf66ffcc5ebb4db..ceb4bab432f184709b570d6b6ff7800121c7fe18 100644
--- a/kernel/trace/ftrace.c
+++ b/kernel/trace/ftrace.c
@@ -1581,12 +1581,7 @@ unsigned long ftrace_location_range(unsigned long start, unsigned long end)
 	key.ip = start;
 	key.flags = end;	/* overload flags, as it is unsigned long */
 
-	/*
-	 * It is in atomic context when called from ftrace_int3_handler(),
-	 * in this case rcu lock is not needed.
-	 */
-	if (!in_atomic())
-		rcu_read_lock();
+	preempt_disable_notrace();
 	for (pg = ftrace_pages_start; pg; pg = pg->next) {
 		if (pg->index == 0 ||
 		    end < pg->records[0].ip ||
@@ -1600,8 +1595,7 @@ unsigned long ftrace_location_range(unsigned long start, unsigned long end)
 			break;
 		}
 	}
-	if (!in_atomic())
-		rcu_read_unlock();
+	preempt_enable_notrace();
 
 	return ip;
 }