return true;
 }
 
+static void kvm_arm_set_pmu(struct kvm *kvm, struct arm_pmu *arm_pmu)
+{
+       lockdep_assert_held(&kvm->arch.config_lock);
+
+       kvm->arch.arm_pmu = arm_pmu;
+}
+
+/**
+ * kvm_arm_set_default_pmu - No PMU set, get the default one.
+ * @kvm: The kvm pointer
+ *
+ * The observant among you will notice that the supported_cpus
+ * mask does not get updated for the default PMU even though it
+ * is quite possible the selected instance supports only a
+ * subset of cores in the system. This is intentional, and
+ * upholds the preexisting behavior on heterogeneous systems
+ * where vCPUs can be scheduled on any core but the guest
+ * counters could stop working.
+ */
+static int kvm_arm_set_default_pmu(struct kvm *kvm)
+{
+       struct arm_pmu *arm_pmu = kvm_pmu_probe_armpmu();
+
+       if (!arm_pmu)
+               return -ENODEV;
+
+       kvm_arm_set_pmu(kvm, arm_pmu);
+       return 0;
+}
+
 static int kvm_arm_pmu_v3_set_pmu(struct kvm_vcpu *vcpu, int pmu_id)
 {
        struct kvm *kvm = vcpu->kvm;
                                break;
                        }
 
-                       kvm->arch.arm_pmu = arm_pmu;
+                       kvm_arm_set_pmu(kvm, arm_pmu);
                        cpumask_copy(kvm->arch.supported_cpus, &arm_pmu->supported_cpus);
                        ret = 0;
                        break;
                return -EBUSY;
 
        if (!kvm->arch.arm_pmu) {
-               /*
-                * No PMU set, get the default one.
-                *
-                * The observant among you will notice that the supported_cpus
-                * mask does not get updated for the default PMU even though it
-                * is quite possible the selected instance supports only a
-                * subset of cores in the system. This is intentional, and
-                * upholds the preexisting behavior on heterogeneous systems
-                * where vCPUs can be scheduled on any core but the guest
-                * counters could stop working.
-                */
-               kvm->arch.arm_pmu = kvm_pmu_probe_armpmu();
-               if (!kvm->arch.arm_pmu)
-                       return -ENODEV;
+               int ret = kvm_arm_set_default_pmu(kvm);
+
+               if (ret)
+                       return ret;
        }
 
        switch (attr->attr) {