GLOBAL_ENTRY(jprobe_inst_return)
        br.call.sptk.many b0=jprobe_break
 END(jprobe_inst_return)
+
+GLOBAL_ENTRY(invalidate_stacked_regs)
+       movl r16=invalidate_restore_cfm
+       ;;
+       mov b6=r16
+       ;;
+       br.ret.sptk.many b6
+       ;;
+invalidate_restore_cfm:
+       mov r16=ar.rsc
+       ;;
+       mov ar.rsc=r0
+       ;;
+       loadrs
+       ;;
+       mov ar.rsc=r16
+       ;;
+       br.cond.sptk.many rp
+END(invalidate_stacked_regs)
+
+GLOBAL_ENTRY(flush_register_stack)
+       // flush dirty regs to backing store (must be first in insn group)
+       flushrs
+       ;;
+       br.ret.sptk.many rp
+END(flush_register_stack)
+
 
        return ret;
 }
 
+struct param_bsp_cfm {
+       unsigned long ip;
+       unsigned long *bsp;
+       unsigned long cfm;
+};
+
+static void ia64_get_bsp_cfm(struct unw_frame_info *info, void *arg)
+{
+       unsigned long ip;
+       struct param_bsp_cfm *lp = arg;
+
+       do {
+               unw_get_ip(info, &ip);
+               if (ip == 0)
+                       break;
+               if (ip == lp->ip) {
+                       unw_get_bsp(info, (unsigned long*)&lp->bsp);
+                       unw_get_cfm(info, (unsigned long*)&lp->cfm);
+                       return;
+               }
+       } while (unw_unwind(info) >= 0);
+       lp->bsp = 0;
+       lp->cfm = 0;
+       return;
+}
+
 int __kprobes setjmp_pre_handler(struct kprobe *p, struct pt_regs *regs)
 {
        struct jprobe *jp = container_of(p, struct jprobe, kp);
        unsigned long addr = ((struct fnptr *)(jp->entry))->ip;
        struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
+       struct param_bsp_cfm pa;
+       int bytes;
+
+       /*
+        * Callee owns the argument space and could overwrite it, eg
+        * tail call optimization. So to be absolutely safe
+        * we save the argument space before transfering the control
+        * to instrumented jprobe function which runs in
+        * the process context
+        */
+       pa.ip = regs->cr_iip;
+       unw_init_running(ia64_get_bsp_cfm, &pa);
+       bytes = (char *)ia64_rse_skip_regs(pa.bsp, pa.cfm & 0x3f)
+                               - (char *)pa.bsp;
+       memcpy( kcb->jprobes_saved_stacked_regs,
+               pa.bsp,
+               bytes );
+       kcb->bsp = pa.bsp;
+       kcb->cfm = pa.cfm;
 
        /* save architectural state */
        kcb->jprobe_saved_regs = *regs;
 int __kprobes longjmp_break_handler(struct kprobe *p, struct pt_regs *regs)
 {
        struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
+       int bytes;
 
+       /* restoring architectural state */
        *regs = kcb->jprobe_saved_regs;
+
+       /* restoring the original argument space */
+       flush_register_stack();
+       bytes = (char *)ia64_rse_skip_regs(kcb->bsp, kcb->cfm & 0x3f)
+                               - (char *)kcb->bsp;
+       memcpy( kcb->bsp,
+               kcb->jprobes_saved_stacked_regs,
+               bytes );
+       invalidate_stacked_regs();
+
        preempt_enable_no_resched();
        return 1;
 }
 
        unsigned long status;
 };
 
+#define        MAX_PARAM_RSE_SIZE      (0x60+0x60/0x3f)
 /* per-cpu kprobe control block */
 struct kprobe_ctlblk {
        unsigned long kprobe_status;
        struct pt_regs jprobe_saved_regs;
+       unsigned long jprobes_saved_stacked_regs[MAX_PARAM_RSE_SIZE];
+       unsigned long *bsp;
+       unsigned long cfm;
        struct prev_kprobe prev_kprobe;
 };
 
 static inline void jprobe_return(void)
 {
 }
+extern void invalidate_stacked_regs(void);
+extern void flush_register_stack(void);
 
 #endif                         /* _ASM_KPROBES_H */