*
  * Is in entry.text as it shouldn't be instrumented.
  */
-SYM_FUNC_START(native_load_gs_index)
+SYM_FUNC_START(asm_load_gs_index)
        FRAME_BEGIN
-       pushfq
-       DISABLE_INTERRUPTS(CLBR_ANY & ~CLBR_RDI)
-       TRACE_IRQS_OFF
        SWAPGS
 .Lgs_change:
        movl    %edi, %gs
 2:     ALTERNATIVE "", "mfence", X86_BUG_SWAPGS_FENCE
        SWAPGS
-       TRACE_IRQS_FLAGS (%rsp)
-       popfq
        FRAME_END
        ret
-SYM_FUNC_END(native_load_gs_index)
-EXPORT_SYMBOL(native_load_gs_index)
+SYM_FUNC_END(asm_load_gs_index)
+EXPORT_SYMBOL(asm_load_gs_index)
 
        _ASM_EXTABLE(.Lgs_change, .Lbad_gs)
        .section .fixup, "ax"
 
 
 #include <asm/nops.h>
 #include <asm/processor-flags.h>
+#include <linux/irqflags.h>
 #include <linux/jump_label.h>
 
 /*
        asm volatile("wbinvd": : :"memory");
 }
 
-extern asmlinkage void native_load_gs_index(unsigned);
+extern asmlinkage void asm_load_gs_index(unsigned int selector);
+
+static inline void native_load_gs_index(unsigned int selector)
+{
+       unsigned long flags;
+
+       local_irq_save(flags);
+       asm_load_gs_index(selector);
+       local_irq_restore(flags);
+}
 
 static inline unsigned long __read_cr4(void)
 {
 
 #ifdef CONFIG_X86_64
 
-static inline void load_gs_index(unsigned selector)
+static inline void load_gs_index(unsigned int selector)
 {
        native_load_gs_index(selector);
 }