/* Restart function and parameter. */
        __u64   restart_fn;                     /* 0x0370 */
        __u64   restart_data;                   /* 0x0378 */
-       __u64   restart_source;                 /* 0x0380 */
+       __u32   restart_source;                 /* 0x0380 */
+       __u32   restart_flags;                  /* 0x0384 */
 
        /* Address space pointer. */
        __u64   kernel_asce;                    /* 0x0388 */
 
 #define _CIF_MCCK_GUEST                BIT(CIF_MCCK_GUEST)
 #define _CIF_DEDICATED_CPU     BIT(CIF_DEDICATED_CPU)
 
+#define RESTART_FLAG_CTLREGS   _AC(1 << 0, U)
+
 #ifndef __ASSEMBLY__
 
 #include <linux/cpumask.h>
 
        OFFSET(__LC_RESTART_FN, lowcore, restart_fn);
        OFFSET(__LC_RESTART_DATA, lowcore, restart_data);
        OFFSET(__LC_RESTART_SOURCE, lowcore, restart_source);
+       OFFSET(__LC_RESTART_FLAGS, lowcore, restart_flags);
        OFFSET(__LC_KERNEL_ASCE, lowcore, kernel_asce);
        OFFSET(__LC_USER_ASCE, lowcore, user_asce);
        OFFSET(__LC_LPP, lowcore, lpp);
 
 4:     j       4b
 ENDPROC(mcck_int_handler)
 
-#
-# PSW restart interrupt handler
-#
 ENTRY(restart_int_handler)
        ALTERNATIVE "", ".insn s,0xb2800000,_LPP_OFFSET", 40
        stg     %r15,__LC_SAVE_AREA_RESTART
+       TSTMSK  __LC_RESTART_FLAGS,RESTART_FLAG_CTLREGS,4
+       jz      0f
+       la      %r15,4095
+       lctlg   %c0,%c15,__LC_CREGS_SAVE_AREA-4095(%r15)
+0:     larl    %r15,.Lstosm_tmp
+       stosm   0(%r15),0x04                    # turn dat on, keep irqs off
        lg      %r15,__LC_RESTART_STACK
        xc      STACK_FRAME_OVERHEAD(__PT_SIZE,%r15),STACK_FRAME_OVERHEAD(%r15)
        stmg    %r0,%r14,STACK_FRAME_OVERHEAD+__PT_R0(%r15)
        xc      0(STACK_FRAME_OVERHEAD,%r15),0(%r15)
        lg      %r1,__LC_RESTART_FN             # load fn, parm & source cpu
        lg      %r2,__LC_RESTART_DATA
-       lg      %r3,__LC_RESTART_SOURCE
+       lgf     %r3,__LC_RESTART_SOURCE
        ltgr    %r3,%r3                         # test source cpu address
        jm      1f                              # negative -> skip source stop
 0:     sigp    %r4,%r3,SIGP_SENSE              # sigp sense to source cpu
 
 
 int diag308(unsigned long subcode, void *addr)
 {
-       if (IS_ENABLED(CONFIG_KASAN))
-               __arch_local_irq_stosm(0x04); /* enable DAT */
        diag_stat_inc(DIAG_STAT_X308);
        return __diag308(subcode, addr);
 }
 
 static void __do_restart(void *ignore)
 {
-       __arch_local_irq_stosm(0x04); /* enable DAT */
        smp_send_stop();
 #ifdef CONFIG_CRASH_DUMP
        crash_kexec(NULL);
 
  */
 static void __machine_kexec(void *data)
 {
-       __arch_local_irq_stosm(0x04); /* enable DAT */
        pfault_fini();
        tracing_off();
        debug_locks_off();
 
        lc->restart_stack = (unsigned long) restart_stack;
        lc->restart_fn = (unsigned long) do_restart;
        lc->restart_data = 0;
-       lc->restart_source = -1UL;
+       lc->restart_source = -1U;
 
        mcck_stack = (unsigned long)memblock_alloc(THREAD_SIZE, THREAD_SIZE);
        if (!mcck_stack)
 
 static void __init setup_lowcore_dat_on(void)
 {
+       struct lowcore *lc = lowcore_ptr[0];
+
        __ctl_clear_bit(0, 28);
        S390_lowcore.external_new_psw.mask |= PSW_MASK_DAT;
        S390_lowcore.svc_new_psw.mask |= PSW_MASK_DAT;
        S390_lowcore.program_new_psw.mask |= PSW_MASK_DAT;
        S390_lowcore.io_new_psw.mask |= PSW_MASK_DAT;
+       __ctl_store(S390_lowcore.cregs_save_area, 0, 15);
        __ctl_set_bit(0, 28);
+       mem_assign_absolute(S390_lowcore.restart_flags, RESTART_FLAG_CTLREGS);
+       mem_assign_absolute(S390_lowcore.program_new_psw, lc->program_new_psw);
+       memcpy_absolute(&S390_lowcore.cregs_save_area, lc->cregs_save_area,
+                       sizeof(S390_lowcore.cregs_save_area));
 }
 
 static struct resource code_resource = {
 
        cpumask_set_cpu(cpu, &init_mm.context.cpu_attach_mask);
        cpumask_set_cpu(cpu, mm_cpumask(&init_mm));
        lc->cpu_nr = cpu;
+       lc->restart_flags = RESTART_FLAG_CTLREGS;
        lc->spinlock_lockval = arch_spin_lockval(cpu);
        lc->spinlock_index = 0;
        lc->percpu_offset = __per_cpu_offset[cpu];
        lc->restart_stack = lc->nodat_stack;
        lc->restart_fn = (unsigned long) func;
        lc->restart_data = (unsigned long) data;
-       lc->restart_source = -1UL;
+       lc->restart_source = -1U;
        pcpu_sigp_retry(pcpu, SIGP_RESTART, 0);
 }
 
        func(data);     /* should not return */
 }
 
-static void __no_sanitize_address pcpu_delegate(struct pcpu *pcpu,
-                                               pcpu_delegate_fn *func,
-                                               void *data, unsigned long stack)
+static void pcpu_delegate(struct pcpu *pcpu,
+                         pcpu_delegate_fn *func,
+                         void *data, unsigned long stack)
 {
        struct lowcore *lc = lowcore_ptr[pcpu - pcpu_devices];
-       unsigned long source_cpu = stap();
+       unsigned int source_cpu = stap();
 
        __load_psw_mask(PSW_KERNEL_BITS | PSW_MASK_DAT);
        if (pcpu->address == source_cpu) {
        __ctl_load(cregs, 0, 15);
 }
 
+static DEFINE_SPINLOCK(ctl_lock);
+static unsigned long ctlreg;
+
 /*
  * Set a bit in a control register of all cpus
  */
 {
        struct ec_creg_mask_parms parms = { 1UL << bit, -1UL, cr };
 
+       spin_lock(&ctl_lock);
+       memcpy_absolute(&ctlreg, &S390_lowcore.cregs_save_area[cr], sizeof(ctlreg));
+       __set_bit(bit, &ctlreg);
+       memcpy_absolute(&S390_lowcore.cregs_save_area[cr], &ctlreg, sizeof(ctlreg));
+       spin_unlock(&ctl_lock);
        on_each_cpu(smp_ctl_bit_callback, &parms, 1);
 }
 EXPORT_SYMBOL(smp_ctl_set_bit);
 {
        struct ec_creg_mask_parms parms = { 0, ~(1UL << bit), cr };
 
+       spin_lock(&ctl_lock);
+       memcpy_absolute(&ctlreg, &S390_lowcore.cregs_save_area[cr], sizeof(ctlreg));
+       __clear_bit(bit, &ctlreg);
+       memcpy_absolute(&S390_lowcore.cregs_save_area[cr], &ctlreg, sizeof(ctlreg));
+       spin_unlock(&ctl_lock);
        on_each_cpu(smp_ctl_bit_callback, &parms, 1);
 }
 EXPORT_SYMBOL(smp_ctl_clear_bit);
 /*
  *     Activate a secondary processor.
  */
-static void __no_sanitize_address smp_start_secondary(void *cpuvoid)
+static void smp_start_secondary(void *cpuvoid)
 {
        S390_lowcore.restart_stack = (unsigned long) restart_stack;
        S390_lowcore.restart_fn = (unsigned long) do_restart;
        S390_lowcore.restart_data = 0;
-       S390_lowcore.restart_source = -1UL;
-       __ctl_load(S390_lowcore.cregs_save_area, 0, 15);
-       __load_psw_mask(PSW_KERNEL_BITS | PSW_MASK_DAT);
+       S390_lowcore.restart_source = -1U;
+       S390_lowcore.restart_flags = 0;
        call_on_stack_noreturn(smp_init_secondary, S390_lowcore.kernel_stack);
 }