/* tracing */
 void _mcount(void);
-unsigned long prepare_ftrace_return(unsigned long parent, unsigned long ip);
+unsigned long prepare_ftrace_return(unsigned long parent, unsigned long ip,
+                                               unsigned long sp);
 
 void pnv_power9_force_smt4_catch(void);
 void pnv_power9_force_smt4_release(void);
 
 #define MCOUNT_ADDR            ((unsigned long)(_mcount))
 #define MCOUNT_INSN_SIZE       4 /* sizeof mcount call */
 
+#define HAVE_FUNCTION_GRAPH_RET_ADDR_PTR
+
 #ifdef __ASSEMBLY__
 
 /* Based off of objdump optput from glibc */
 
                 * FIXME: IMHO these tests do not belong in
                 * arch-dependent code, they are generic.
                 */
-               ip = ftrace_graph_ret_addr(tsk, &graph_idx, ip, NULL);
+               ip = ftrace_graph_ret_addr(tsk, &graph_idx, ip, stack);
 #ifdef CONFIG_KPROBES
                /*
                 * Mark stacktraces with kretprobed functions on them
 
  * Hook the return address and push it in the stack of return addrs
  * in current thread info. Return the address we want to divert to.
  */
-unsigned long prepare_ftrace_return(unsigned long parent, unsigned long ip)
+unsigned long prepare_ftrace_return(unsigned long parent, unsigned long ip,
+                                               unsigned long sp)
 {
        unsigned long return_hooker;
 
 
        return_hooker = ppc_function_entry(return_to_handler);
 
-       if (!function_graph_enter(parent, ip, 0, NULL))
+       if (!function_graph_enter(parent, ip, 0, (unsigned long *)sp))
                parent = return_hooker;
 out:
        return parent;
 
 
 #ifdef CONFIG_FUNCTION_GRAPH_TRACER
 _GLOBAL(ftrace_graph_caller)
+       addi    r5, r1, 48
        /* load r4 with local address */
        lwz     r4, 44(r1)
        subi    r4, r4, MCOUNT_INSN_SIZE
 
        std     r2, 24(r1)
        ld      r2, PACATOC(r13)        /* get kernel TOC in r2 */
 
+       addi    r5, r1, 112
        mfctr   r4              /* ftrace_caller has moved local addr here */
        std     r4, 40(r1)
        mflr    r3              /* ftrace_caller has restored LR from stack */
 
 
 #ifdef CONFIG_FUNCTION_GRAPH_TRACER
 _GLOBAL(ftrace_graph_caller)
+       addi    r5, r1, 112
        /* load r4 with local address */
        ld      r4, 128(r1)
        subi    r4, r4, MCOUNT_INSN_SIZE