return 0;
 }
 
+#ifdef CONFIG_MPROFILE_KERNEL
+
+#define PACATOC offsetof(struct paca_struct, kernel_toc)
+
+/*
+ * ld      r12,PACATOC(r13)
+ * addis   r12,r12,<high>
+ * addi    r12,r12,<low>
+ * mtctr   r12
+ * bctr
+ */
+static u32 stub_insns[] = {
+       PPC_INST_LD | __PPC_RT(R12) | __PPC_RA(R13) | PACATOC,
+       PPC_INST_ADDIS | __PPC_RT(R12) | __PPC_RA(R12),
+       PPC_INST_ADDI | __PPC_RT(R12) | __PPC_RA(R12),
+       PPC_INST_MTCTR | __PPC_RS(R12),
+       PPC_INST_BCTR,
+};
+
+/*
+ * For mprofile-kernel we use a special stub for ftrace_caller() because we
+ * can't rely on r2 containing this module's TOC when we enter the stub.
+ *
+ * That can happen if the function calling us didn't need to use the toc. In
+ * that case it won't have setup r2, and the r2 value will be either the
+ * kernel's toc, or possibly another modules toc.
+ *
+ * To deal with that this stub uses the kernel toc, which is always accessible
+ * via the paca (in r13). The target (ftrace_caller()) is responsible for
+ * saving and restoring the toc before returning.
+ */
+static inline int create_ftrace_stub(struct ppc64_stub_entry *entry,
+                                       unsigned long addr,
+                                       struct module *me)
+{
+       long reladdr;
+
+       memcpy(entry->jump, stub_insns, sizeof(stub_insns));
+
+       /* Stub uses address relative to kernel toc (from the paca) */
+       reladdr = addr - kernel_toc_addr();
+       if (reladdr > 0x7FFFFFFF || reladdr < -(0x80000000L)) {
+               pr_err("%s: Address of %ps out of range of kernel_toc.\n",
+                                                       me->name, (void *)addr);
+               return 0;
+       }
+
+       entry->jump[1] |= PPC_HA(reladdr);
+       entry->jump[2] |= PPC_LO(reladdr);
+
+       /* Eventhough we don't use funcdata in the stub, it's needed elsewhere. */
+       entry->funcdata = func_desc(addr);
+       entry->magic = STUB_MAGIC;
+
+       return 1;
+}
+
+static bool is_mprofile_ftrace_call(const char *name)
+{
+       if (!strcmp("_mcount", name))
+               return true;
+#ifdef CONFIG_DYNAMIC_FTRACE
+       if (!strcmp("ftrace_caller", name))
+               return true;
+#ifdef CONFIG_DYNAMIC_FTRACE_WITH_REGS
+       if (!strcmp("ftrace_regs_caller", name))
+               return true;
+#endif
+#endif
+
+       return false;
+}
+#else
+static inline int create_ftrace_stub(struct ppc64_stub_entry *entry,
+                                       unsigned long addr,
+                                       struct module *me)
+{
+       return 0;
+}
+
+static bool is_mprofile_ftrace_call(const char *name)
+{
+       return false;
+}
+#endif
+
 /*
  * r2 is the TOC pointer: it actually points 0x8000 into the TOC (this gives the
  * value maximum span in an instruction which uses a signed offset). Round down
 static inline int create_stub(const Elf64_Shdr *sechdrs,
                              struct ppc64_stub_entry *entry,
                              unsigned long addr,
-                             struct module *me)
+                             struct module *me,
+                             const char *name)
 {
        long reladdr;
 
+       if (is_mprofile_ftrace_call(name))
+               return create_ftrace_stub(entry, addr, me);
+
        memcpy(entry->jump, ppc64_stub_insns, sizeof(ppc64_stub_insns));
 
        /* Stub uses address relative to r2. */
    stub to set up the TOC ptr (r2) for the function. */
 static unsigned long stub_for_addr(const Elf64_Shdr *sechdrs,
                                   unsigned long addr,
-                                  struct module *me)
+                                  struct module *me,
+                                  const char *name)
 {
        struct ppc64_stub_entry *stubs;
        unsigned int i, num_stubs;
                        return (unsigned long)&stubs[i];
        }
 
-       if (!create_stub(sechdrs, &stubs[i], addr, me))
+       if (!create_stub(sechdrs, &stubs[i], addr, me, name))
                return 0;
 
        return (unsigned long)&stubs[i];
 }
 
-#ifdef CONFIG_MPROFILE_KERNEL
-static bool is_mprofile_ftrace_call(const char *name)
-{
-       if (!strcmp("_mcount", name))
-               return true;
-
-       return false;
-}
-
-/*
- * In case of _mcount calls, do not save the current callee's TOC (in r2) into
- * the original caller's stack frame. If we did we would clobber the saved TOC
- * value of the original caller.
- */
-static void squash_toc_save_inst(const char *name, unsigned long addr)
-{
-       struct ppc64_stub_entry *stub = (struct ppc64_stub_entry *)addr;
-
-       /* Only for calls to _mcount */
-       if (strcmp("_mcount", name) != 0)
-               return;
-
-       stub->jump[2] = PPC_INST_NOP;
-}
-#else
-static void squash_toc_save_inst(const char *name, unsigned long addr) { }
-
-static bool is_mprofile_ftrace_call(const char *name)
-{
-       return false;
-}
-#endif
-
 /* We expect a noop next: if it is, replace it with instruction to
    restore r2. */
 static int restore_r2(const char *name, u32 *instruction, struct module *me)
                        if (sym->st_shndx == SHN_UNDEF ||
                            sym->st_shndx == SHN_LIVEPATCH) {
                                /* External: go via stub */
-                               value = stub_for_addr(sechdrs, value, me);
+                               value = stub_for_addr(sechdrs, value, me,
+                                               strtab + sym->st_name);
                                if (!value)
                                        return -ENOENT;
                                if (!restore_r2(strtab + sym->st_name,
                                                        (u32 *)location + 1, me))
                                        return -ENOEXEC;
-
-                               squash_toc_save_inst(strtab + sym->st_name, value);
                        } else
                                value += local_entry_offset(sym);
 
        return 0;
 }
 
-#ifdef CONFIG_MPROFILE_KERNEL
-
-#define PACATOC offsetof(struct paca_struct, kernel_toc)
-
-/*
- * For mprofile-kernel we use a special stub for ftrace_caller() because we
- * can't rely on r2 containing this module's TOC when we enter the stub.
- *
- * That can happen if the function calling us didn't need to use the toc. In
- * that case it won't have setup r2, and the r2 value will be either the
- * kernel's toc, or possibly another modules toc.
- *
- * To deal with that this stub uses the kernel toc, which is always accessible
- * via the paca (in r13). The target (ftrace_caller()) is responsible for
- * saving and restoring the toc before returning.
- */
-static unsigned long create_ftrace_stub(const Elf64_Shdr *sechdrs,
-                               struct module *me, unsigned long addr)
-{
-       struct ppc64_stub_entry *entry;
-       unsigned int i, num_stubs;
-       /*
-        * ld      r12,PACATOC(r13)
-        * addis   r12,r12,<high>
-        * addi    r12,r12,<low>
-        * mtctr   r12
-        * bctr
-        */
-       static u32 stub_insns[] = {
-               PPC_INST_LD | __PPC_RT(R12) | __PPC_RA(R13) | PACATOC,
-               PPC_INST_ADDIS | __PPC_RT(R12) | __PPC_RA(R12),
-               PPC_INST_ADDI | __PPC_RT(R12) | __PPC_RA(R12),
-               PPC_INST_MTCTR | __PPC_RS(R12),
-               PPC_INST_BCTR,
-       };
-       long reladdr;
-
-       num_stubs = sechdrs[me->arch.stubs_section].sh_size / sizeof(*entry);
-
-       /* Find the next available stub entry */
-       entry = (void *)sechdrs[me->arch.stubs_section].sh_addr;
-       for (i = 0; i < num_stubs && stub_func_addr(entry->funcdata); i++, entry++);
-
-       if (i >= num_stubs) {
-               pr_err("%s: Unable to find a free slot for ftrace stub.\n", me->name);
-               return 0;
-       }
-
-       memcpy(entry->jump, stub_insns, sizeof(stub_insns));
-
-       /* Stub uses address relative to kernel toc (from the paca) */
-       reladdr = addr - kernel_toc_addr();
-       if (reladdr > 0x7FFFFFFF || reladdr < -(0x80000000L)) {
-               pr_err("%s: Address of %ps out of range of kernel_toc.\n",
-                                                       me->name, (void *)addr);
-               return 0;
-       }
-
-       entry->jump[1] |= PPC_HA(reladdr);
-       entry->jump[2] |= PPC_LO(reladdr);
-
-       /* Eventhough we don't use funcdata in the stub, it's needed elsewhere. */
-       entry->funcdata = func_desc(addr);
-       entry->magic = STUB_MAGIC;
-
-       return (unsigned long)entry;
-}
-#else
-static unsigned long create_ftrace_stub(const Elf64_Shdr *sechdrs,
-                               struct module *me, unsigned long addr)
-{
-       return stub_for_addr(sechdrs, addr, me);
-}
-#endif
-
 int module_finalize_ftrace(struct module *mod, const Elf_Shdr *sechdrs)
 {
-       mod->arch.tramp = create_ftrace_stub(sechdrs, mod,
-                                       (unsigned long)ftrace_caller);
+       mod->arch.tramp = stub_for_addr(sechdrs,
+                                       (unsigned long)ftrace_caller,
+                                       mod,
+                                       "ftrace_caller");
 #ifdef CONFIG_DYNAMIC_FTRACE_WITH_REGS
-       mod->arch.tramp_regs = create_ftrace_stub(sechdrs, mod,
-                                       (unsigned long)ftrace_regs_caller);
+       mod->arch.tramp_regs = stub_for_addr(sechdrs,
+                                       (unsigned long)ftrace_regs_caller,
+                                       mod,
+                                       "ftrace_regs_caller");
        if (!mod->arch.tramp_regs)
                return -ENOENT;
 #endif