#define BRANCH_SET_LINK        0x1
 #define BRANCH_ABSOLUTE        0x2
 
+DECLARE_STATIC_KEY_FALSE(init_mem_is_free);
+
 bool is_offset_in_branch_range(long offset);
 bool is_offset_in_cond_branch_range(long offset);
 int create_branch(ppc_inst_t *instr, const u32 *addr,
 
 #include <linux/init.h>
 #include <linux/cpuhotplug.h>
 #include <linux/uaccess.h>
+#include <linux/jump_label.h>
 
 #include <asm/tlbflush.h>
 #include <asm/page.h>
 
 #endif /* CONFIG_STRICT_KERNEL_RWX */
 
+__ro_after_init DEFINE_STATIC_KEY_FALSE(init_mem_is_free);
+
 int patch_instruction(u32 *addr, ppc_inst_t instr)
 {
        /* Make sure we aren't patching a freed init section */
-       if (system_state >= SYSTEM_FREEING_INITMEM && init_section_contains(addr, 4))
+       if (static_branch_likely(&init_mem_is_free) && init_section_contains(addr, 4))
                return 0;
 
        return do_patch_instruction(addr, instr);
 
 #include <asm/kasan.h>
 #include <asm/svm.h>
 #include <asm/mmzone.h>
+#include <asm/code-patching.h>
 
 #include <mm/mmu_decl.h>
 
 {
        ppc_md.progress = ppc_printk_progress;
        mark_initmem_nx();
+       static_branch_enable(&init_mem_is_free);
        free_initmem_default(POISON_FREE_INITMEM);
 }