if (copy_from_user(>od, (void __user *)attr->addr, sizeof(gtod)))
                return -EFAULT;
 
-       if (test_kvm_facility(kvm, 139))
-               kvm_s390_set_tod_clock_ext(kvm, >od);
-       else if (gtod.epoch_idx == 0)
-               kvm_s390_set_tod_clock(kvm, gtod.tod);
-       else
+       if (!test_kvm_facility(kvm, 139) && gtod.epoch_idx)
                return -EINVAL;
+       kvm_s390_set_tod_clock(kvm, >od);
 
        VM_EVENT(kvm, 3, "SET: TOD extension: 0x%x, TOD base: 0x%llx",
                gtod.epoch_idx, gtod.tod);
 
 static int kvm_s390_set_tod_low(struct kvm *kvm, struct kvm_device_attr *attr)
 {
-       u64 gtod;
+       struct kvm_s390_vm_tod_clock gtod = { 0 };
 
-       if (copy_from_user(>od, (void __user *)attr->addr, sizeof(gtod)))
+       if (copy_from_user(>od.tod, (void __user *)attr->addr,
+                          sizeof(gtod.tod)))
                return -EFAULT;
 
-       kvm_s390_set_tod_clock(kvm, gtod);
-       VM_EVENT(kvm, 3, "SET: TOD base: 0x%llx", gtod);
+       kvm_s390_set_tod_clock(kvm, >od);
+       VM_EVENT(kvm, 3, "SET: TOD base: 0x%llx", gtod.tod);
        return 0;
 }
 
        return 0;
 }
 
-void kvm_s390_set_tod_clock_ext(struct kvm *kvm,
-                                const struct kvm_s390_vm_tod_clock *gtod)
+void kvm_s390_set_tod_clock(struct kvm *kvm,
+                           const struct kvm_s390_vm_tod_clock *gtod)
 {
        struct kvm_vcpu *vcpu;
        struct kvm_s390_tod_clock_ext htod;
        get_tod_clock_ext((char *)&htod);
 
        kvm->arch.epoch = gtod->tod - htod.tod;
-       kvm->arch.epdx = gtod->epoch_idx - htod.epoch_idx;
-
-       if (kvm->arch.epoch > gtod->tod)
-               kvm->arch.epdx -= 1;
+       kvm->arch.epdx = 0;
+       if (test_kvm_facility(kvm, 139)) {
+               kvm->arch.epdx = gtod->epoch_idx - htod.epoch_idx;
+               if (kvm->arch.epoch > gtod->tod)
+                       kvm->arch.epdx -= 1;
+       }
 
        kvm_s390_vcpu_block_all(kvm);
        kvm_for_each_vcpu(i, vcpu, kvm) {
        mutex_unlock(&kvm->lock);
 }
 
-void kvm_s390_set_tod_clock(struct kvm *kvm, u64 tod)
-{
-       struct kvm_vcpu *vcpu;
-       int i;
-
-       mutex_lock(&kvm->lock);
-       preempt_disable();
-       kvm->arch.epoch = tod - get_tod_clock();
-       kvm_s390_vcpu_block_all(kvm);
-       kvm_for_each_vcpu(i, vcpu, kvm)
-               vcpu->arch.sie_block->epoch = kvm->arch.epoch;
-       kvm_s390_vcpu_unblock_all(kvm);
-       preempt_enable();
-       mutex_unlock(&kvm->lock);
-}
-
 /**
  * kvm_arch_fault_in_page - fault-in guest page if necessary
  * @vcpu: The corresponding virtual cpu
 
 int kvm_s390_handle_sigp_pei(struct kvm_vcpu *vcpu);
 
 /* implemented in kvm-s390.c */
-void kvm_s390_set_tod_clock_ext(struct kvm *kvm,
-                                const struct kvm_s390_vm_tod_clock *gtod);
-void kvm_s390_set_tod_clock(struct kvm *kvm, u64 tod);
+void kvm_s390_set_tod_clock(struct kvm *kvm,
+                           const struct kvm_s390_vm_tod_clock *gtod);
 long kvm_arch_fault_in_page(struct kvm_vcpu *vcpu, gpa_t gpa, int writable);
 int kvm_s390_store_status_unloaded(struct kvm_vcpu *vcpu, unsigned long addr);
 int kvm_s390_vcpu_store_status(struct kvm_vcpu *vcpu, unsigned long addr);
 
 /* Handle SCK (SET CLOCK) interception */
 static int handle_set_clock(struct kvm_vcpu *vcpu)
 {
+       struct kvm_s390_vm_tod_clock gtod = { 0 };
        int rc;
        u8 ar;
-       u64 op2, val;
+       u64 op2;
 
        vcpu->stat.instruction_sck++;
 
        op2 = kvm_s390_get_base_disp_s(vcpu, &ar);
        if (op2 & 7)    /* Operand must be on a doubleword boundary */
                return kvm_s390_inject_program_int(vcpu, PGM_SPECIFICATION);
-       rc = read_guest(vcpu, op2, ar, &val, sizeof(val));
+       rc = read_guest(vcpu, op2, ar, >od.tod, sizeof(gtod.tod));
        if (rc)
                return kvm_s390_inject_prog_cond(vcpu, rc);
 
-       VCPU_EVENT(vcpu, 3, "SCK: setting guest TOD to 0x%llx", val);
-       kvm_s390_set_tod_clock(vcpu->kvm, val);
+       VCPU_EVENT(vcpu, 3, "SCK: setting guest TOD to 0x%llx", gtod.tod);
+       kvm_s390_set_tod_clock(vcpu->kvm, >od);
 
        kvm_s390_set_psw_cc(vcpu, 0);
        return 0;