struct dmub_psr *psr = dc->res_pool->psr;
 
        if ((psr != NULL) && link->psr_feature_enabled)
-               psr->funcs->set_psr_enable(psr, allow_active);
+               psr->funcs->psr_enable(psr, allow_active);
        else if ((dmcu != NULL && dmcu->funcs->is_dmcu_initialized(dmcu)) && link->psr_feature_enabled)
                dmcu->funcs->set_psr_enable(dmcu, allow_active, wait);
 
        struct dmub_psr *psr = dc->res_pool->psr;
 
        if (psr != NULL && link->psr_feature_enabled)
-               psr->funcs->get_psr_state(psr_state);
+               psr->funcs->psr_get_state(psr_state);
        else if (dmcu != NULL && link->psr_feature_enabled)
                dmcu->funcs->get_psr_state(dmcu, psr_state);
 
        psr_context->frame_delay = 0;
 
        if (psr)
-               link->psr_feature_enabled = psr->funcs->setup_psr(psr, link, psr_context);
+               link->psr_feature_enabled = psr->funcs->psr_copy_settings(psr, link, psr_context);
        else
                link->psr_feature_enabled = dmcu->funcs->setup_psr(dmcu, link, psr_context);
 
 
 /**
  * Get PSR state from firmware.
  */
-static void dmub_get_psr_state(uint32_t *psr_state)
+static void dmub_psr_get_state(uint32_t *psr_state)
 {
        // Not yet implemented
        // Trigger GPINT interrupt from firmware
 }
 
+static void dmub_psr_set_version(struct dmub_psr *dmub, struct dc_stream_state *stream)
+{
+       //stream->psr_version;
+       union dmub_rb_cmd cmd;
+       struct dc_context *dc = dmub->ctx;
+
+       cmd.psr_set_version.header.type = DMUB_CMD__PSR;
+       cmd.psr_set_version.header.sub_type = DMUB_CMD__PSR_SET_VERSION;
+
+       if (stream->psr_version == 0x0)
+               return;
+       else if (stream->psr_version == 0x1)
+               cmd.psr_set_version.psr_set_version_data.version = PSR_VERSION_1;
+       else if (stream->psr_version == 0x2)
+               cmd.psr_set_version.psr_set_version_data.version = PSR_VERSION_2;
+
+       cmd.psr_enable.header.payload_bytes = sizeof(struct dmub_cmd_psr_set_version_data);
+
+       dc_dmub_srv_cmd_queue(dc->dmub_srv, &cmd.psr_enable.header);
+       dc_dmub_srv_cmd_execute(dc->dmub_srv);
+       dc_dmub_srv_wait_idle(dc->dmub_srv);
+}
+
 /**
  * Enable/Disable PSR.
  */
-static void dmub_set_psr_enable(struct dmub_psr *dmub, bool enable)
+static void dmub_psr_enable(struct dmub_psr *dmub, bool enable)
 {
        union dmub_rb_cmd cmd;
        struct dc_context *dc = dmub->ctx;
 /**
  * Set PSR level.
  */
-static void dmub_set_psr_level(struct dmub_psr *dmub, uint16_t psr_level)
+static void dmub_psr_set_level(struct dmub_psr *dmub, uint16_t psr_level)
 {
        union dmub_rb_cmd cmd;
        uint32_t psr_state = 0;
        struct dc_context *dc = dmub->ctx;
 
-       dmub_get_psr_state(&psr_state);
+       dmub_psr_get_state(&psr_state);
 
        if (psr_state == 0)
                return;
 /**
  * Setup PSR by programming phy registers and sending psr hw context values to firmware.
  */
-static bool dmub_setup_psr(struct dmub_psr *dmub,
+static bool dmub_psr_copy_settings(struct dmub_psr *dmub,
                struct dc_link *link,
                struct psr_context *psr_context)
 {
 
        for (int i = 0; i < MAX_PIPES; i++) {
                if (res_ctx &&
-                               res_ctx->pipe_ctx[i].stream &&
-                               res_ctx->pipe_ctx[i].stream->link &&
-                               res_ctx->pipe_ctx[i].stream->link == link &&
-                               res_ctx->pipe_ctx[i].stream->link->connector_signal == SIGNAL_TYPE_EDP) {
+                       res_ctx->pipe_ctx[i].stream &&
+                       res_ctx->pipe_ctx[i].stream->link &&
+                       res_ctx->pipe_ctx[i].stream->link == link &&
+                       res_ctx->pipe_ctx[i].stream->link->connector_signal == SIGNAL_TYPE_EDP) {
                        pipe_ctx = &res_ctx->pipe_ctx[i];
                        break;
                }
        }
 
-       if (!pipe_ctx ||
-                       !&pipe_ctx->plane_res ||
-                       !&pipe_ctx->stream_res)
+       if (!pipe_ctx || !&pipe_ctx->plane_res || !&pipe_ctx->stream_res)
                return false;
 
        // Program DP DPHY fast training registers
        copy_settings_data->hyst_lines                          = psr_context->hyst_lines;
        copy_settings_data->phy_type                            = psr_context->phyType;
        copy_settings_data->aux_repeat                          = psr_context->aux_repeats;
-       copy_settings_data->smu_optimizations_en        = psr_context->allow_smu_optimizations;
-       copy_settings_data->skip_wait_for_pll_lock      = psr_context->skipPsrWaitForPllLock;
+       copy_settings_data->smu_optimizations_en                = psr_context->allow_smu_optimizations;
+       copy_settings_data->skip_wait_for_pll_lock              = psr_context->skipPsrWaitForPllLock;
        copy_settings_data->frame_delay                         = psr_context->frame_delay;
        copy_settings_data->smu_phy_id                          = psr_context->smuPhyId;
-       copy_settings_data->num_of_controllers          = psr_context->numberOfControllers;
+       copy_settings_data->num_of_controllers                  = psr_context->numberOfControllers;
        copy_settings_data->frame_cap_ind                       = psr_context->psrFrameCaptureIndicationReq;
-       copy_settings_data->phy_num                                     = psr_context->frame_delay & 0x7;
+       copy_settings_data->phy_num                             = psr_context->frame_delay & 0x7;
        copy_settings_data->link_rate                           = psr_context->frame_delay & 0xF;
 
        dc_dmub_srv_cmd_queue(dc->dmub_srv, &cmd.psr_copy_settings.header);
 }
 
 static const struct dmub_psr_funcs psr_funcs = {
-       .set_psr_enable                 = dmub_set_psr_enable,
-       .setup_psr                              = dmub_setup_psr,
-       .get_psr_state                  = dmub_get_psr_state,
-       .set_psr_level                  = dmub_set_psr_level,
+       .psr_set_version                = dmub_psr_set_version,
+       .psr_copy_settings              = dmub_psr_copy_settings,
+       .psr_enable                     = dmub_psr_enable,
+       .psr_get_state                  = dmub_psr_get_state,
+       .psr_set_level                  = dmub_psr_set_level,
 };
 
 /**
 
        struct dmub_cmd_header header;
 };
 
-struct dmub_cmd_psr_setup_data {
+struct dmub_cmd_psr_set_version_data {
        enum psr_version version; // PSR version 1 or 2
 };
 
-struct dmub_rb_cmd_psr_setup {
+struct dmub_rb_cmd_psr_set_version {
        struct dmub_cmd_header header;
-       struct dmub_cmd_psr_setup_data psr_setup_data;
+       struct dmub_cmd_psr_set_version_data psr_set_version_data;
 };
 
 union dmub_rb_cmd {
        struct dmub_rb_cmd_enable_disp_power_gating enable_disp_power_gating;
        struct dmub_rb_cmd_dpphy_init dpphy_init;
        struct dmub_rb_cmd_dig1_transmitter_control dig1_transmitter_control;
-       struct dmub_rb_cmd_psr_enable psr_enable;
+       struct dmub_rb_cmd_psr_set_version psr_set_version;
        struct dmub_rb_cmd_psr_copy_settings psr_copy_settings;
+       struct dmub_rb_cmd_psr_enable psr_enable;
        struct dmub_rb_cmd_psr_set_level psr_set_level;
        struct dmub_rb_cmd_PLAT_54186_wa PLAT_54186_wa;
-       struct dmub_rb_cmd_psr_setup psr_setup;
 };
 
 #pragma pack(pop)