int kvm_vcpu_ioctl_dirty_tlb(struct kvm_vcpu *vcpu,
                             struct kvm_dirty_tlb *cfg);
 
+long kvmppc_alloc_lpid(void);
+void kvmppc_claim_lpid(long lpid);
+void kvmppc_free_lpid(long lpid);
+void kvmppc_init_lpid(unsigned long nr_lpids);
+
 #endif /* __POWERPC_KVM_PPC_H__ */
 
 
 /* POWER7 has 10-bit LPIDs, PPC970 has 6-bit LPIDs */
 #define MAX_LPID_970   63
-#define NR_LPIDS       (LPID_RSVD + 1)
-unsigned long lpid_inuse[BITS_TO_LONGS(NR_LPIDS)];
 
 long kvmppc_alloc_hpt(struct kvm *kvm)
 {
        unsigned long hpt;
-       unsigned long lpid;
+       long lpid;
        struct revmap_entry *rev;
        struct kvmppc_linear_info *li;
 
        }
        kvm->arch.revmap = rev;
 
-       /* Allocate the guest's logical partition ID */
-       do {
-               lpid = find_first_zero_bit(lpid_inuse, NR_LPIDS);
-               if (lpid >= NR_LPIDS) {
-                       pr_err("kvm_alloc_hpt: No LPIDs free\n");
-                       goto out_freeboth;
-               }
-       } while (test_and_set_bit(lpid, lpid_inuse));
+       lpid = kvmppc_alloc_lpid();
+       if (lpid < 0)
+               goto out_freeboth;
 
        kvm->arch.sdr1 = __pa(hpt) | (HPT_ORDER - 18);
        kvm->arch.lpid = lpid;
 
 void kvmppc_free_hpt(struct kvm *kvm)
 {
-       clear_bit(kvm->arch.lpid, lpid_inuse);
+       kvmppc_free_lpid(kvm->arch.lpid);
        vfree(kvm->arch.revmap);
        if (kvm->arch.hpt_li)
                kvm_release_hpt(kvm->arch.hpt_li);
        if (!cpu_has_feature(CPU_FTR_HVMODE))
                return -EINVAL;
 
-       memset(lpid_inuse, 0, sizeof(lpid_inuse));
-
+       /* POWER7 has 10-bit LPIDs, PPC970 and e500mc have 6-bit LPIDs */
        if (cpu_has_feature(CPU_FTR_ARCH_206)) {
                host_lpid = mfspr(SPRN_LPID);   /* POWER7 */
                rsvd_lpid = LPID_RSVD;
                rsvd_lpid = MAX_LPID_970;
        }
 
-       set_bit(host_lpid, lpid_inuse);
+       kvmppc_init_lpid(rsvd_lpid + 1);
+
+       kvmppc_claim_lpid(host_lpid);
        /* rsvd_lpid is reserved for use in partition switching */
-       set_bit(rsvd_lpid, lpid_inuse);
+       kvmppc_claim_lpid(rsvd_lpid);
 
        return 0;
 }
 
        return r;
 }
 
+static unsigned long lpid_inuse[BITS_TO_LONGS(KVMPPC_NR_LPIDS)];
+static unsigned long nr_lpids;
+
+long kvmppc_alloc_lpid(void)
+{
+       long lpid;
+
+       do {
+               lpid = find_first_zero_bit(lpid_inuse, KVMPPC_NR_LPIDS);
+               if (lpid >= nr_lpids) {
+                       pr_err("%s: No LPIDs free\n", __func__);
+                       return -ENOMEM;
+               }
+       } while (test_and_set_bit(lpid, lpid_inuse));
+
+       return lpid;
+}
+
+void kvmppc_claim_lpid(long lpid)
+{
+       set_bit(lpid, lpid_inuse);
+}
+
+void kvmppc_free_lpid(long lpid)
+{
+       clear_bit(lpid, lpid_inuse);
+}
+
+void kvmppc_init_lpid(unsigned long nr_lpids_param)
+{
+       nr_lpids = min_t(unsigned long, KVMPPC_NR_LPIDS, nr_lpids_param);
+       memset(lpid_inuse, 0, sizeof(lpid_inuse));
+}
+
 int kvm_arch_init(void *opaque)
 {
        return 0;