}
        }
 
-       ret = 0;
-       if (enable_remapped_mode) {
-               /* Use legacy mode in IRTE */
-               struct amd_iommu_pi_data pi;
-
-               /**
-                * Here, pi is used to:
-                * - Tell IOMMU to use legacy mode for this interrupt.
-                */
-               pi.is_guest_mode = false;
-               ret = irq_set_vcpu_affinity(host_irq, &pi);
-       }
+       if (enable_remapped_mode)
+               ret = irq_set_vcpu_affinity(host_irq, NULL);
+       else
+               ret = 0;
 out:
        srcu_read_unlock(&kvm->irq_srcu, idx);
        return ret;
 
 {
        int ret;
        struct amd_iommu_pi_data *pi_data = vcpu_info;
-       struct vcpu_data *vcpu_pi_info = pi_data->vcpu_data;
        struct amd_ir_data *ir_data = data->chip_data;
        struct irq_2_irte *irte_info = &ir_data->irq_2_irte;
        struct iommu_dev_data *dev_data;
                return -EINVAL;
 
        ir_data->cfg = irqd_cfg(data);
-       pi_data->ir_data = ir_data;
 
-       if (pi_data->is_guest_mode) {
+       if (pi_data) {
+               struct vcpu_data *vcpu_pi_info = pi_data->vcpu_data;
+
+               pi_data->ir_data = ir_data;
+
+               WARN_ON_ONCE(!pi_data->is_guest_mode);
+
                ir_data->ga_root_ptr = (vcpu_pi_info->pi_desc_addr >> 12);
                ir_data->ga_vector = vcpu_pi_info->vector;
                ir_data->ga_tag = pi_data->ga_tag;