diff --git a/xen/arch/x86/domctl.c b/xen/arch/x86/domctl.c
index 3d7fc34..8e5502a 100644
--- a/xen/arch/x86/domctl.c
+++ b/xen/arch/x86/domctl.c
@@ -1386,21 +1386,21 @@ long arch_do_domctl(
         switch ( domctl->u.psr_cat_op.cmd )
         {
         case XEN_DOMCTL_PSR_CAT_OP_SET_L3_CBM:
-            ret = psr_set_l3_cbm(d, domctl->u.psr_cat_op.target,
-                                 domctl->u.psr_cat_op.data,
-                                 PSR_CBM_TYPE_L3);
+            ret = psr_set_val(d, domctl->u.psr_cat_op.target,
+                              domctl->u.psr_cat_op.data,
+                              PSR_CBM_TYPE_L3);
             break;
 
         case XEN_DOMCTL_PSR_CAT_OP_SET_L3_CODE:
-            ret = psr_set_l3_cbm(d, domctl->u.psr_cat_op.target,
-                                 domctl->u.psr_cat_op.data,
-                                 PSR_CBM_TYPE_L3_CODE);
+            ret = psr_set_val(d, domctl->u.psr_cat_op.target,
+                              domctl->u.psr_cat_op.data,
+                              PSR_CBM_TYPE_L3_CODE);
             break;
 
         case XEN_DOMCTL_PSR_CAT_OP_SET_L3_DATA:
-            ret = psr_set_l3_cbm(d, domctl->u.psr_cat_op.target,
-                                 domctl->u.psr_cat_op.data,
-                                 PSR_CBM_TYPE_L3_DATA);
+            ret = psr_set_val(d, domctl->u.psr_cat_op.target,
+                              domctl->u.psr_cat_op.data,
+                              PSR_CBM_TYPE_L3_DATA);
             break;
 
         case XEN_DOMCTL_PSR_CAT_OP_GET_L3_CBM:
diff --git a/xen/arch/x86/psr.c b/xen/arch/x86/psr.c
index fdf4e3a..32f899a 100644
--- a/xen/arch/x86/psr.c
+++ b/xen/arch/x86/psr.c
@@ -39,6 +39,17 @@ enum psr_feat_type {
 struct feat_node;
 struct psr_cat_socket_info;
 
+/*
+ * To add a new PSR feature, please follow below steps:
+ * 1. Implement feature specific callback functions listed in
+ *    "struct feat_ops";
+ * 2. Implement a "struct feat_ops" instance for the feature and register
+ *    the feature callback functions into it;
+ * 3. Delcare feat_node instance for the feature and malloc memory for it
+ *    in cat_cpu_prepare(). Correspondingly, free it in free_feature();
+ * 4. Add feature initialization codes in cat_cpu_init().
+ */
+
 /* Every feature enabled MUST implement such ops and callback functions. */
 struct feat_ops {
     /*
@@ -57,6 +68,74 @@ struct feat_ops {
                    enum cbm_type type, uint64_t *val);
     /* get_max_cos_max is used to get feature's cos_max. */
     unsigned int (*get_max_cos_max)(const struct feat_node *feat);
+    /*
+     * get_cos_num is used to get the COS registers amount used by the
+     * feature for one setting, e.g. CDP uses 2 COSs but CAT uses 1.
+     */
+    unsigned int (*get_cos_num)(const struct feat_node *feat);
+    /*
+     * get_old_val and set_new_val are a pair of functions called together.
+     * The caller will traverse all features in the list and call both
+     * functions for every feature to do below two things:
+     * 1. get old_cos register value of all supported features and
+     * 2. set the new value for the feature.
+     *
+     * All the values are set into value array according the traversal order,
+     * meaning the same order of feature list members.
+     *
+     * The return value is the amount of entries to skip in the value array
+     * or error.
+     * 1 - one entry in value array.
+     * 2 - two entries in value array, e.g. CDP uses two entries.
+     * negative - error.
+     */
+    int (*get_old_val)(uint64_t val[],
+                       const struct feat_node *feat,
+                       unsigned int old_cos);
+    int (*set_new_val)(uint64_t val[],
+                       const struct feat_node *feat,
+                       unsigned int old_cos,
+                       enum cbm_type type,
+                       uint64_t m);
+    /*
+     * compare_val is used in set value process to compare if the
+     * input value array can match all the features' COS registers values
+     * according to input cos id.
+     *
+     * The return value is the amount of entries to skip in the value array
+     * or error.
+     * 1 - one entry in value array.
+     * 2 - two entries in value array, e.g. CDP uses two entries.
+     * negative - error.
+     */
+    int (*compare_val)(const uint64_t val[], const struct feat_node *feat,
+                        unsigned int cos, bool *found);
+    /*
+     * get_cos_max_from_type is used to get the cos_max value of the feature
+     * according to input type.
+     */
+    unsigned int (*get_cos_max_from_type)(const struct feat_node *feat,
+                                          enum cbm_type type);
+    /*
+     * exceeds_cos_max is used to check if the input cos id exceeds the
+     * feature's cos_max and if the input value is not the default one.
+     * Even if the associated cos exceeds the cos_max, HW can work with default
+     * value. That is the reason we need check if input value is default one.
+     * If both criteria are fulfilled, that means the input exceeds the
+     * range.
+     *
+     * The return value of the function means the number of the value array
+     * entries to skip or error.
+     * 1 - one entry in value array.
+     * 2 - two entries in value array, e.g. CDP uses two entries.
+     * 0 - error, exceed cos_max and the input value is not default.
+     */
+    unsigned int (*exceeds_cos_max)(const uint64_t val[],
+                                    const struct feat_node *feat,
+                                    unsigned int cos);
+    /* write_msr is used to write out feature MSR register. */
+    int (*write_msr)(unsigned int cos, const uint64_t val[],
+                     struct feat_node *feat);
 };
 
 /* CAT/CDP HW info data structure. */
@@ -128,18 +207,7 @@ static DEFINE_PER_CPU(struct psr_assoc, psr_assoc);
 /* Feature list entry of feature L3 CAT/CDP. */
 static struct feat_node *feat_l3;
 
-static struct feat_node *get_feat_l3(struct psr_cat_socket_info *info)
-{
-    struct feat_node *feat_tmp;
-
-    list_for_each_entry(feat_tmp, &info->feat_list, list)
-        if ( feat_tmp->feature == PSR_SOCKET_L3_CAT ||
-             feat_tmp->feature == PSR_SOCKET_L3_CDP )
-            return feat_tmp;
-
-    return NULL;
-}
-
+/* Common functions. */
 static void free_feature(struct psr_cat_socket_info *info)
 {
     struct feat_node *feat_tmp;
@@ -162,6 +230,29 @@ static void free_feature(struct psr_cat_socket_info *info)
     }
 }
 
+static bool_t psr_check_cbm(unsigned int cbm_len, uint64_t cbm)
+{
+    unsigned int first_bit, zero_bit;
+
+    /* Set bits should only in the range of [0, cbm_len). */
+    if ( cbm & (~0ull << cbm_len) )
+        return 0;
+
+    /* At least one bit need to be set. */
+    if ( cbm == 0 )
+        return 0;
+
+    first_bit = find_first_bit(&cbm, cbm_len);
+    zero_bit = find_next_zero_bit(&cbm, cbm_len, first_bit);
+
+    /* Set bits should be contiguous. */
+    if ( zero_bit < cbm_len &&
+         find_next_bit(&cbm, cbm_len, zero_bit) < cbm_len )
+        return 0;
+
+    return 1;
+}
+
 /* L3 CAT/CDP callback functions implementation. */
 static void l3_cat_init_feature(unsigned int eax, unsigned int ebx,
                                 unsigned int ecx, unsigned int edx,
@@ -286,11 +377,235 @@ static unsigned int l3_cat_get_max_cos_max(const struct feat_node *feat)
     return feat->info.cos_max;
 }
 
+static unsigned int l3_cat_get_cos_num(const struct feat_node *feat)
+{
+    if ( feat->feature == PSR_SOCKET_L3_CDP )
+        return 2;
+
+    return 1;
+}
+
+static int l3_cat_get_old_val(uint64_t val[],
+                              const struct feat_node *feat,
+                              unsigned int old_cos)
+{
+    if ( old_cos > feat->info.cos_max )
+        /* Use default value. */
+        old_cos = 0;
+
+    /* CDP */
+    if ( feat->feature == PSR_SOCKET_L3_CDP )
+    {
+        /* Data */
+        val[0] = get_cdp_data(feat, old_cos);
+        /* Code */
+        val[1] = get_cdp_code(feat, old_cos);
+
+        /* CDP uses two COSs, one for data, one for code. */
+        return 2;
+    }
+
+    /* CAT */
+    val[0] =  feat->cos_reg_val[old_cos];
+
+    /* L3 CAT uses one COS. */
+    return 1;
+}
+
+static int l3_cat_set_new_val(uint64_t val[],
+                              const struct feat_node *feat,
+                              unsigned int old_cos,
+                              enum cbm_type type,
+                              uint64_t m)
+{
+    /* L3 CAT uses one COS. */
+    unsigned int ret = 1;
+
+    switch ( type ) {
+    /* CDP uses two COS, one for data, one for code. */
+    case PSR_CBM_TYPE_L3_DATA:
+    case PSR_CBM_TYPE_L3_CODE:
+        if ( feat->feature != PSR_SOCKET_L3_CDP )
+            return -ENXIO;
+
+        if ( !psr_check_cbm(feat->info.cbm_len, m) )
+            return -EINVAL;
+
+        if ( type == PSR_CBM_TYPE_L3_DATA )
+            val[0] = m;
+        else
+            val[1] = m;
+
+        ret= 2;
+        break;
+
+    /* CAT */
+    case PSR_CBM_TYPE_L3:
+        if ( !psr_check_cbm(feat->info.cbm_len, m) )
+            return -EINVAL;
+
+        val[0] = m;
+        break;
+
+    default:
+        if ( feat->feature == PSR_SOCKET_L3_CDP )
+            /* CDP uses two COS, one for data, one for code. */
+            ret = 2;
+        break;
+    }
+
+    return ret;
+}
+
+static int l3_cat_compare_val(const uint64_t val[],
+                              const struct feat_node *feat,
+                              unsigned int cos, bool *found)
+{
+    uint64_t l3_def_cbm;
+
+    l3_def_cbm = (1ull << feat->info.cbm_len) - 1;
+
+    /* CDP */
+    if ( feat->feature == PSR_SOCKET_L3_CDP )
+    {
+        if ( cos > feat->info.cos_max )
+        {
+            if ( val[0] != l3_def_cbm ||
+                 val[1] != l3_def_cbm )
+            {
+                *found = false;
+                return -ENOENT;
+            }
+            *found = true;
+        }
+        else
+            *found = (val[0] == get_cdp_data(feat, cos) &&
+                      val[1] == get_cdp_code(feat, cos));
+
+        /* CDP uses two COS, one for data, one for code. */
+        return 2;
+    }
+
+    /* CAT */
+    if ( cos > feat->info.cos_max )
+    {
+        if ( val[0] != l3_def_cbm )
+        {
+            *found = false;
+            return -ENOENT;
+        }
+        *found = true;
+    }
+    else
+        *found = (val[0] == feat->cos_reg_val[cos]);
+
+    /* L3 CAT uses one COS. */
+    return 1;
+}
+
+static unsigned int l3_cat_get_cos_max_from_type(const struct feat_node *feat,
+                                                 enum cbm_type type)
+{
+    if ( type != PSR_CBM_TYPE_L3_DATA &&
+         type != PSR_CBM_TYPE_L3_CODE &&
+         type != PSR_CBM_TYPE_L3 )
+        return 0;
+
+    return feat->info.cos_max;
+}
+
+static unsigned int l3_cat_exceeds_cos_max(const uint64_t val[],
+                                           const struct feat_node *feat,
+                                           unsigned int cos)
+{
+    uint64_t l3_def_cbm;
+
+    l3_def_cbm = (1ull << feat->info.cbm_len) - 1;
+
+    /* CDP */
+    if ( feat->feature == PSR_SOCKET_L3_CDP )
+    {
+        if ( cos > feat->info.cos_max &&
+             (val[0] != l3_def_cbm || val[1] != l3_def_cbm) )
+                /*
+                 * Exceed cos_max and value to set is not default,
+                 * return error.
+                 */
+                return 0;
+
+        /* CDP uses two COS, one for data, one for code. */
+        return 2;
+    }
+
+    /* CAT */
+    if ( cos > feat->info.cos_max &&
+         val[0] != l3_def_cbm )
+            /*
+             * Exceed cos_max and value to set is not default,
+             * return error.
+             */
+            return 0;
+
+    /* L3 CAT uses one COS. */
+    return 1;
+}
+
+static int l3_cat_write_msr(unsigned int cos, const uint64_t val[],
+                            struct feat_node *feat)
+{
+    /* CDP */
+    if ( feat->feature == PSR_SOCKET_L3_CDP )
+    {
+        /*
+         * If input cos is more than the cos_max of the feature, we should
+         * not set the value.
+         */
+        if ( cos > feat->info.cos_max )
+            /* CDP uses two COS, one for data, one for code. */
+            return 2;
+
+        /* Data */
+        if ( get_cdp_data(feat, cos) != val[0] )
+        {
+            get_cdp_data(feat, cos) = val[0];
+            wrmsrl(MSR_IA32_PSR_L3_MASK_DATA(cos), val[0]);
+        }
+        /* Code */
+        if ( get_cdp_code(feat, cos) != val[1] )
+        {
+            get_cdp_code(feat, cos) = val[1];
+            wrmsrl(MSR_IA32_PSR_L3_MASK_CODE(cos), val[1]);
+        }
+        /* CDP uses two COS, one for data, one for code. */
+        return 2;
+    }
+
+    /* CAT */
+    if ( cos > feat->info.cos_max )
+        /* L3 CAT uses one COS. */
+        return 1;
+
+    if ( feat->cos_reg_val[cos] != val[0] )
+    {
+        feat->cos_reg_val[cos] = val[0];
+        wrmsrl(MSR_IA32_PSR_L3_MASK(cos), val[0]);
+    }
+    /* L3 CAT uses one COS. */
+    return 1;
+}
+
 struct feat_ops l3_cat_ops = {
     .init_feature = l3_cat_init_feature,
     .get_feat_info = l3_cat_get_feat_info,
     .get_val = l3_cat_get_val,
     .get_max_cos_max = l3_cat_get_max_cos_max,
+    .get_cos_num = l3_cat_get_cos_num,
+    .get_old_val = l3_cat_get_old_val,
+    .set_new_val = l3_cat_set_new_val,
+    .compare_val = l3_cat_compare_val,
+    .get_cos_max_from_type = l3_cat_get_cos_max_from_type,
+    .exceeds_cos_max = l3_cat_exceeds_cos_max,
+    .write_msr = l3_cat_write_msr,
 };
 
 static unsigned int get_socket_cpu(unsigned int socket)
@@ -516,11 +831,6 @@ static struct psr_cat_socket_info *get_cat_socket_info(unsigned int socket)
     return cat_socket_info + socket;
 }
 
-static inline bool_t cdp_is_enabled(unsigned int socket)
-{
-    return test_bit(PSR_SOCKET_L3_CDP, &cat_socket_info[socket].feat_mask);
-}
-
 int psr_get_info(unsigned int socket, enum cbm_type type,
                  uint32_t dat[], uint32_t array_len)
 {
@@ -562,194 +872,314 @@ int psr_get_val(struct domain *d, unsigned int socket,
     return 0;
 }
 
-static bool_t psr_check_cbm(unsigned int cbm_len, uint64_t cbm)
-{
-    unsigned int first_bit, zero_bit;
-
-    /* Set bits should only in the range of [0, cbm_len). */
-    if ( cbm & (~0ull << cbm_len) )
-        return 0;
-
-    /* At least one bit need to be set. */
-    if ( cbm == 0 )
-        return 0;
-
-    first_bit = find_first_bit(&cbm, cbm_len);
-    zero_bit = find_next_zero_bit(&cbm, cbm_len, first_bit);
-
-    /* Set bits should be contiguous. */
-    if ( zero_bit < cbm_len &&
-         find_next_bit(&cbm, cbm_len, zero_bit) < cbm_len )
-        return 0;
-
-    return 1;
-}
-
-struct cos_cbm_info
+struct cos_write_info
 {
     unsigned int cos;
-    bool_t cdp;
-    uint64_t cbm_code;
-    uint64_t cbm_data;
+    struct list_head *feat_list;
+    const uint64_t *val;
 };
 
-static void do_write_l3_cbm(void *data)
+static void do_write_psr_msr(void *data)
 {
-    struct cos_cbm_info *info = data;
+    struct cos_write_info *info = (struct cos_write_info *)data;
+    unsigned int cos           = info->cos;
+    struct list_head *feat_list= info->feat_list;
+    const uint64_t *val        = info->val;
+    struct feat_node *feat_tmp;
+    int ret;
+
+    if ( !feat_list )
+        return;
 
-    if ( info->cdp )
+    list_for_each_entry(feat_tmp, feat_list, list)
     {
-        wrmsrl(MSR_IA32_PSR_L3_MASK_CODE(info->cos), info->cbm_code);
-        wrmsrl(MSR_IA32_PSR_L3_MASK_DATA(info->cos), info->cbm_data);
+        ret = feat_tmp->ops.write_msr(cos, val, feat_tmp);
+        if ( ret <= 0)
+            return;
+
+        val += ret;
     }
-    else
-        wrmsrl(MSR_IA32_PSR_L3_MASK(info->cos), info->cbm_code);
 }
 
-static int write_l3_cbm(unsigned int socket, unsigned int cos,
-                        uint64_t cbm_code, uint64_t cbm_data, bool_t cdp)
+static int write_psr_msr(unsigned int socket, unsigned int cos,
+                         const uint64_t *val)
 {
-    struct cos_cbm_info info =
+    struct psr_cat_socket_info *info = get_cat_socket_info(socket);
+
+    struct cos_write_info data =
     {
         .cos = cos,
-        .cbm_code = cbm_code,
-        .cbm_data = cbm_data,
-        .cdp = cdp,
+        .feat_list = &info->feat_list,
+        .val = val,
     };
 
     if ( socket == cpu_to_socket(smp_processor_id()) )
-        do_write_l3_cbm(&info);
+        do_write_psr_msr(&data);
     else
     {
         unsigned int cpu = get_socket_cpu(socket);
 
         if ( cpu >= nr_cpu_ids )
             return -ENOTSOCK;
-        on_selected_cpus(cpumask_of(cpu), do_write_l3_cbm, &info, 1);
+        on_selected_cpus(cpumask_of(cpu), do_write_psr_msr, &data, 1);
     }
 
     return 0;
 }
 
-static int find_cos(struct feat_node *feat, unsigned int *ref,
-                    unsigned int cos_max,
-                    uint64_t cbm_code, uint64_t cbm_data, bool_t cdp_enabled)
+static unsigned int get_cos_num(const struct psr_cat_socket_info *info)
+{
+    const struct feat_node *feat_tmp;
+    unsigned int num = 0;
+
+    list_for_each_entry(feat_tmp, &info->feat_list, list)
+        num += feat_tmp->ops.get_cos_num(feat_tmp);
+
+    return num;
+}
+
+static int get_old_set_new(uint64_t *val,
+                           uint32_t array_len,
+                           const struct psr_cat_socket_info *info,
+                           unsigned int old_cos,
+                           enum cbm_type type,
+                           uint64_t m)
+{
+    const struct feat_node *feat_tmp;
+    int ret;
+    uint64_t *val_tmp = val;
+
+    if ( !val )
+        return -EINVAL;
+
+    /* Get all features current values according to old_cos. */
+    list_for_each_entry(feat_tmp, &info->feat_list, list)
+    {
+        /* value getting order is same as feature list */
+        ret = feat_tmp->ops.get_old_val(val_tmp, feat_tmp, old_cos);
+
+        val_tmp += ret;
+        if ( val_tmp - val > array_len)
+            return -EINVAL;
+    }
+
+    /* Set new value into array according to feature's position in array. */
+    val_tmp = val;
+    list_for_each_entry(feat_tmp, &info->feat_list, list)
+    {
+        /* value setting order is same as feature list */
+        ret = feat_tmp->ops.set_new_val(val_tmp, feat_tmp, old_cos, type, m);
+        if ( ret < 0 )
+            return ret;
+
+        val_tmp += ret;
+        if ( val_tmp - val > array_len)
+            return -EINVAL;
+    }
+
+    return 0;
+}
+
+static int find_cos(const uint64_t *val, uint32_t array_len,
+                    enum cbm_type type,
+                    const struct psr_cat_socket_info *info)
 {
     unsigned int cos;
+    const unsigned int *ref = info->cos_ref;
+    const struct feat_node *feat_tmp;
+    const uint64_t *val_tmp = val;
+    int ret;
+    bool found = false;
+    unsigned int cos_max = 0;
+
+    /* cos_max is the one of the feature which is being set. */
+    list_for_each_entry(feat_tmp, &info->feat_list, list)
+    {
+        cos_max = feat_tmp->ops.get_cos_max_from_type(feat_tmp, type);
+        if ( cos_max > 0 )
+            break;
+    }
 
     for ( cos = 0; cos <= cos_max; cos++ )
     {
-        if ( (ref[cos] || cos == 0) &&
-             ((!cdp_enabled && feat->cos_reg_val[cos] == cbm_code) ||
-              (cdp_enabled && get_cdp_code(feat, cos) == cbm_code &&
-                              get_cdp_data(feat, cos) == cbm_data)) )
+        if ( cos && !ref[cos] )
+            continue;
+
+        /* Not found, need find again from beginning. */
+        val_tmp = val;
+        list_for_each_entry(feat_tmp, &info->feat_list, list)
+        {
+            /*
+             * Compare value according to feature list order.
+             * We must follow this order because value array is assembled
+             * as this order in get_old_set_new().
+             */
+            ret = feat_tmp->ops.compare_val(val_tmp, feat_tmp, cos, &found);
+            if ( ret < 0 )
+                return ret;
+
+            /* If fail to match, go to next cos to compare. */
+            if ( !found )
+                break;
+
+            val_tmp += ret;
+            if ( val_tmp - val > array_len )
+                return -EINVAL;
+        }
+
+        /*
+         * With this cos id, every entry of value array can match. This cos
+         * is what we find.
+         */
+        if ( found )
             return cos;
     }
 
     return -ENOENT;
 }
 
-static int pick_avail_cos(unsigned int *ref, unsigned int cos_max,
-                          unsigned int old_cos)
+static bool exceeds_cos_max(const uint64_t *val,
+                            uint32_t array_len,
+                            const struct psr_cat_socket_info *info,
+                            unsigned int cos)
+{
+    unsigned int ret;
+    const uint64_t *val_tmp = val;
+    const struct feat_node *feat_tmp;
+
+    list_for_each_entry(feat_tmp, &info->feat_list, list)
+    {
+        ret = feat_tmp->ops.exceeds_cos_max(val_tmp, feat_tmp, cos);
+        if ( !ret )
+            return false;
+
+        val_tmp += ret;
+        if ( val_tmp - val > array_len )
+            return false;
+    }
+
+    return true;
+}
+
+static int alloc_new_cos(const struct psr_cat_socket_info *info,
+                         const uint64_t *val, uint32_t array_len,
+                         unsigned int old_cos,
+                         enum cbm_type type)
 {
     unsigned int cos;
+    unsigned int cos_max = 0;
+    const struct feat_node *feat_tmp;
+    const unsigned int *ref = info->cos_ref;
+
+    /*
+     * cos_max is the one of the feature which is being set.
+     */
+    list_for_each_entry(feat_tmp, &info->feat_list, list)
+    {
+        cos_max = feat_tmp->ops.get_cos_max_from_type(feat_tmp, type);
+        if ( cos_max > 0 )
+            break;
+    }
+
+    if ( !cos_max )
+        return -ENOENT;
 
     /* If old cos is referred only by the domain, then use it. */
-    if ( ref[old_cos] == 1 && old_cos != 0 )
-        return old_cos;
+    if ( ref[old_cos] == 1 && old_cos )
+        if ( exceeds_cos_max(val, array_len, info, old_cos) )
+            return old_cos;
 
     /* Find an unused one other than cos0. */
     for ( cos = 1; cos <= cos_max; cos++ )
-        if ( ref[cos] == 0 )
+        /*
+         * ref is 0 means this COS is not used by other domain and
+         * can be used for current setting.
+         */
+        if ( !ref[cos] )
+        {
+            if ( !exceeds_cos_max(val, array_len, info, cos) )
+                return -ENOENT;
+
             return cos;
+        }
 
     return -ENOENT;
 }
 
-int psr_set_l3_cbm(struct domain *d, unsigned int socket,
-                   uint64_t cbm, enum cbm_type type)
+int psr_set_val(struct domain *d, unsigned int socket,
+                uint64_t val, enum cbm_type type)
 {
-    unsigned int old_cos, cos_max;
+    unsigned int old_cos;
     int cos, ret;
-    uint64_t cbm_data, cbm_code;
-    bool_t cdp_enabled = cdp_is_enabled(socket);
-    struct psr_cat_socket_info *info = get_cat_socket_info(socket);
     unsigned int *ref;
-    struct feat_node *feat_tmp;
+    uint64_t *val_array;
+    struct psr_cat_socket_info *info = get_cat_socket_info(socket);
+    uint32_t array_len;
 
     if ( IS_ERR(info) )
         return PTR_ERR(info);
 
-    feat_tmp = get_feat_l3(info);
-    if ( !feat_tmp )
-        return -ENOENT;
-
-    if ( !psr_check_cbm(feat_tmp->info.cbm_len, cbm) )
-        return -EINVAL;
-
-    if ( !cdp_enabled && (type == PSR_CBM_TYPE_L3_CODE ||
-                          type == PSR_CBM_TYPE_L3_DATA) )
-        return -ENXIO;
-
-    cos_max = feat_tmp->info.cos_max;
     old_cos = d->arch.psr_cos_ids[socket];
     ref = info->cos_ref;
 
-    switch ( type )
-    {
-    case PSR_CBM_TYPE_L3:
-        cbm_code = cbm;
-        cbm_data = cbm;
-        break;
-
-    case PSR_CBM_TYPE_L3_CODE:
-        cbm_code = cbm;
-        cbm_data = get_cdp_data(feat_tmp, old_cos);
-        break;
-
-    case PSR_CBM_TYPE_L3_DATA:
-        cbm_code = get_cdp_code(feat_tmp, old_cos);
-        cbm_data = cbm;
-        break;
+    /*
+     * Assemle a value array to store all featues cos_reg_val[old_cos].
+     * And, set the input val into array according to the feature's
+     * position in array.
+     */
+    array_len = get_cos_num((const struct psr_cat_socket_info *)info);
+    val_array = xzalloc_array(uint64_t, array_len);
+    if ( !val_array )
+        return -ENOMEM;
 
-    default:
-        ASSERT_UNREACHABLE();
-        return -EINVAL;
+    if ( (ret = get_old_set_new(val_array, array_len,
+                                (const struct psr_cat_socket_info *)info,
+                                old_cos, type, val)) != 0 )
+    {
+        xfree(val_array);
+        return ret;
     }
 
+    /*
+     * Lock here to make sure the ref is not changed during find and
+     * write process.
+     */
     spin_lock(&info->ref_lock);
-    cos = find_cos(feat_tmp, ref, cos_max, cbm_code, cbm_data, cdp_enabled);
+    /*
+     * Try to find if there is already a COS ID on which all features' values
+     * are same as the array. Then, we can reuse this COS ID.
+     */
+    cos = find_cos((const uint64_t *)val_array, array_len, type,
+                   (const struct psr_cat_socket_info *)info);
     if ( cos >= 0 )
     {
         if ( cos == old_cos )
         {
             spin_unlock(&info->ref_lock);
+            xfree(val_array);
             return 0;
         }
     }
     else
     {
-        cos = pick_avail_cos(ref, cos_max, old_cos);
+        /* If fail to find, we need allocate a new COS ID. */
+        cos = alloc_new_cos((const struct psr_cat_socket_info *)info,
+                            (const uint64_t *)val_array, array_len,
+                            old_cos, type);
         if ( cos < 0 )
         {
             spin_unlock(&info->ref_lock);
+            xfree(val_array);
             return cos;
         }
 
-        /* We try to avoid writing MSR. */
-        if ( (cdp_enabled &&
-             (get_cdp_code(feat_tmp, cos) != cbm_code ||
-              get_cdp_data(feat_tmp, cos) != cbm_data)) ||
-             (!cdp_enabled && feat_tmp->cos_reg_val[cos] != cbm_code) )
+        /* Write all features MSRs according to the COS ID. */
+        ret = write_psr_msr(socket, cos, (const uint64_t *)val_array);
+        if ( ret )
         {
-            ret = write_l3_cbm(socket, cos, cbm_code, cbm_data, cdp_enabled);
-            if ( ret )
-            {
-                spin_unlock(&info->ref_lock);
-                return ret;
-            }
-            get_cdp_code(feat_tmp, cos) = cbm_code;
-            get_cdp_data(feat_tmp, cos) = cbm_data;
+            spin_unlock(&info->ref_lock);
+            xfree(val_array);
+            return ret;
         }
     }
 
@@ -758,6 +1188,7 @@ int psr_set_l3_cbm(struct domain *d, unsigned int socket,
     spin_unlock(&info->ref_lock);
 
     d->arch.psr_cos_ids[socket] = cos;
+    xfree(val_array);
 
     return 0;
 }
diff --git a/xen/include/asm-x86/psr.h b/xen/include/asm-x86/psr.h
index f2872a2..17ee6f3 100644
--- a/xen/include/asm-x86/psr.h
+++ b/xen/include/asm-x86/psr.h
@@ -67,8 +67,8 @@ int psr_get_info(unsigned int socket, enum cbm_type type,
                  uint32_t dat[], uint32_t array_len);
 int psr_get_val(struct domain *d, unsigned int socket,
                 uint64_t *val, enum cbm_type type);
-int psr_set_l3_cbm(struct domain *d, unsigned int socket,
-                   uint64_t cbm, enum cbm_type type);
+int psr_set_val(struct domain *d, unsigned int socket,
+                uint64_t val, enum cbm_type type);
 
 int psr_domain_init(struct domain *d);
 void psr_domain_free(struct domain *d);
