diff --git a/xen/arch/x86/psr.c b/xen/arch/x86/psr.c
index 0b5073c..ee20389 100644
--- a/xen/arch/x86/psr.c
+++ b/xen/arch/x86/psr.c
@@ -23,6 +23,9 @@
 #define PSR_CAT        (1<<1)
 #define PSR_CDP        (1<<2)
 
+/* Per spec, the maximum COS register number is 128. */
+#define MAX_COS_REG_NUM  128
+
 struct psr_cat_cbm {
     union {
         uint64_t cbm;
@@ -31,14 +34,19 @@ struct psr_cat_cbm {
             uint64_t data;
         };
     };
-    unsigned int ref;
 };
 
 struct psr_cat_socket_info {
     unsigned int cbm_len;
     unsigned int cos_max;
     struct psr_cat_cbm *cos_to_cbm;
-    spinlock_t cbm_lock;
+    /*
+     * Every entry of cos_ref is the reference count of a COS register.
+     * One entry of cos_ref corresponds to one COS ID.
+     */
+    unsigned int cos_ref[MAX_COS_REG_NUM];
+    /* Protect cos_ref */
+    spinlock_t ref_lock;
 };
 
 struct psr_assoc {
@@ -54,7 +62,7 @@ static unsigned long *__read_mostly cdp_socket_enable;
 
 static unsigned int opt_psr;
 static unsigned int __initdata opt_rmid_max = 255;
-static unsigned int __read_mostly opt_cos_max = 255;
+static unsigned int __read_mostly opt_cos_max = MAX_COS_REG_NUM - 1;
 static uint64_t rmid_mask;
 static DEFINE_PER_CPU(struct psr_assoc, psr_assoc);
 
@@ -400,14 +408,15 @@ static int write_l3_cbm(unsigned int socket, unsigned int cos,
     return 0;
 }
 
-static int find_cos(struct psr_cat_cbm *map, unsigned int cos_max,
+static int find_cos(struct psr_cat_cbm *map, unsigned int *ref,
+                    unsigned int cos_max,
                     uint64_t cbm_code, uint64_t cbm_data, bool_t cdp_enabled)
 {
     unsigned int cos;
 
     for ( cos = 0; cos <= cos_max; cos++ )
     {
-        if ( (map[cos].ref || cos == 0) &&
+        if ( (ref[cos] || cos == 0) &&
              ((!cdp_enabled && map[cos].cbm == cbm_code) ||
               (cdp_enabled && map[cos].code == cbm_code &&
                               map[cos].data == cbm_data)) )
@@ -417,18 +426,18 @@ static int find_cos(struct psr_cat_cbm *map, unsigned int cos_max,
     return -ENOENT;
 }
 
-static int pick_avail_cos(struct psr_cat_cbm *map, unsigned int cos_max,
+static int pick_avail_cos(unsigned int *ref, unsigned int cos_max,
                           unsigned int old_cos)
 {
     unsigned int cos;
 
     /* If old cos is referred only by the domain, then use it. */
-    if ( map[old_cos].ref == 1 && old_cos != 0 )
+    if ( ref[old_cos] == 1 && old_cos != 0 )
         return old_cos;
 
     /* Find an unused one other than cos0. */
     for ( cos = 1; cos <= cos_max; cos++ )
-        if ( map[cos].ref == 0 )
+        if ( ref[cos] == 0 )
             return cos;
 
     return -ENOENT;
@@ -443,6 +452,7 @@ int psr_set_l3_cbm(struct domain *d, unsigned int socket,
     bool_t cdp_enabled = cdp_is_enabled(socket);
     struct psr_cat_cbm *map;
     struct psr_cat_socket_info *info = get_cat_socket_info(socket);
+    unsigned int *ref;
 
     if ( IS_ERR(info) )
         return PTR_ERR(info);
@@ -457,6 +467,7 @@ int psr_set_l3_cbm(struct domain *d, unsigned int socket,
     cos_max = info->cos_max;
     old_cos = d->arch.psr_cos_ids[socket];
     map = info->cos_to_cbm;
+    ref = info->cos_ref;
 
     switch ( type )
     {
@@ -480,22 +491,22 @@ int psr_set_l3_cbm(struct domain *d, unsigned int socket,
         return -EINVAL;
     }
 
-    spin_lock(&info->cbm_lock);
-    cos = find_cos(map, cos_max, cbm_code, cbm_data, cdp_enabled);
+    spin_lock(&info->ref_lock);
+    cos = find_cos(map, ref, cos_max, cbm_code, cbm_data, cdp_enabled);
     if ( cos >= 0 )
     {
         if ( cos == old_cos )
         {
-            spin_unlock(&info->cbm_lock);
+            spin_unlock(&info->ref_lock);
             return 0;
         }
     }
     else
     {
-        cos = pick_avail_cos(map, cos_max, old_cos);
+        cos = pick_avail_cos(ref, cos_max, old_cos);
         if ( cos < 0 )
         {
-            spin_unlock(&info->cbm_lock);
+            spin_unlock(&info->ref_lock);
             return cos;
         }
 
@@ -507,7 +518,7 @@ int psr_set_l3_cbm(struct domain *d, unsigned int socket,
             ret = write_l3_cbm(socket, cos, cbm_code, cbm_data, cdp_enabled);
             if ( ret )
             {
-                spin_unlock(&info->cbm_lock);
+                spin_unlock(&info->ref_lock);
                 return ret;
             }
             map[cos].code = cbm_code;
@@ -515,9 +526,9 @@ int psr_set_l3_cbm(struct domain *d, unsigned int socket,
         }
     }
 
-    map[cos].ref++;
-    map[old_cos].ref--;
-    spin_unlock(&info->cbm_lock);
+    ref[cos]++;
+    ref[old_cos]--;
+    spin_unlock(&info->ref_lock);
 
     d->arch.psr_cos_ids[socket] = cos;
 
@@ -540,9 +551,9 @@ static void psr_free_cos(struct domain *d)
             continue;
 
         info = cat_socket_info + socket;
-        spin_lock(&info->cbm_lock);
-        info->cos_to_cbm[cos].ref--;
-        spin_unlock(&info->cbm_lock);
+        spin_lock(&info->ref_lock);
+        info->cos_ref[cos]--;
+        spin_unlock(&info->ref_lock);
     }
 
     xfree(d->arch.psr_cos_ids);
@@ -574,7 +585,7 @@ static int cat_cpu_prepare(unsigned int cpu)
 
     if ( temp_cos_to_cbm == NULL &&
          (temp_cos_to_cbm = xzalloc_array(struct psr_cat_cbm,
-                                          opt_cos_max + 1UL)) == NULL )
+                                          MAX_COS_REG_NUM)) == NULL )
         return -ENOMEM;
 
     return 0;
@@ -609,7 +620,7 @@ static void cat_cpu_init(void)
         /* cos=0 is reserved as default cbm(all ones). */
         info->cos_to_cbm[0].cbm = (1ull << info->cbm_len) - 1;
 
-        spin_lock_init(&info->cbm_lock);
+        spin_lock_init(&info->ref_lock);
 
         set_bit(socket, cat_socket_enable);
 
