}
 }
 
+static void rv1_enable_pme_wa(struct clk_mgr *clk_mgr_base)
+{
+       struct clk_mgr_internal *clk_mgr = TO_CLK_MGR_INTERNAL(clk_mgr_base);
+       struct pp_smu_funcs_rv *pp_smu = NULL;
+
+       if (clk_mgr->pp_smu) {
+               pp_smu = &clk_mgr->pp_smu->rv_funcs;
+
+               if (pp_smu->set_pme_wa_enable)
+                       pp_smu->set_pme_wa_enable(&pp_smu->pp_smu);
+       }
+}
+
 static struct clk_mgr_funcs rv1_clk_funcs = {
        .get_dp_ref_clk_frequency = dce12_get_dp_ref_freq_khz,
        .update_clocks = rv1_update_clocks,
+       .enable_pme_wa = rv1_enable_pme_wa,
 };
 
 static struct clk_mgr_internal_funcs rv1_clk_internal_funcs = {
 
        memset(&(clk_mgr->clks), 0, sizeof(struct dc_clocks));
 }
 
+void dcn2_enable_pme_wa(struct clk_mgr *clk_mgr_base)
+{
+       struct clk_mgr_internal *clk_mgr = TO_CLK_MGR_INTERNAL(clk_mgr_base);
+       struct pp_smu_funcs_nv *pp_smu = NULL;
+
+       if (clk_mgr->pp_smu) {
+               pp_smu = &clk_mgr->pp_smu->nv_funcs;
+
+               if (pp_smu->set_pme_wa_enable)
+                       pp_smu->set_pme_wa_enable(&pp_smu->pp_smu);
+       }
+}
+
 static struct clk_mgr_funcs dcn2_funcs = {
        .get_dp_ref_clk_frequency = dce12_get_dp_ref_freq_khz,
        .update_clocks = dcn2_update_clocks,
-       .init_clocks = dcn2_init_clocks
+       .init_clocks = dcn2_init_clocks,
+       .enable_pme_wa = dcn2_enable_pme_wa
 };
 
 
 
                edp_receiver_ready_T9(link);
 }
 
-// Static helper function which calls the correct function
-// based on pp_smu version
-static void set_pme_wa_enable_by_version(struct dc *dc)
-{
-       struct pp_smu_funcs *pp_smu = NULL;
-
-       if (dc->res_pool->pp_smu)
-               pp_smu = dc->res_pool->pp_smu;
-
-       if (pp_smu) {
-               if (pp_smu->ctx.ver == PP_SMU_VER_RV && pp_smu->rv_funcs.set_pme_wa_enable)
-                       pp_smu->rv_funcs.set_pme_wa_enable(&(pp_smu->rv_funcs.pp_smu));
-#ifdef CONFIG_DRM_AMD_DC_DCN2_0
-               else if (pp_smu->ctx.ver == PP_SMU_VER_NV && pp_smu->nv_funcs.set_pme_wa_enable)
-                       pp_smu->nv_funcs.set_pme_wa_enable(&(pp_smu->nv_funcs.pp_smu));
-#endif
-       }
-}
-
 void dce110_enable_audio_stream(struct pipe_ctx *pipe_ctx)
 {
        /* notify audio driver for audio modes of monitor */
        struct dc *core_dc = pipe_ctx->stream->ctx->dc;
        struct pp_smu_funcs *pp_smu = NULL;
+       struct clk_mgr *clk_mgr = core_dc->clk_mgr;
        unsigned int i, num_audio = 1;
 
        if (pipe_ctx->stream_res.audio && pipe_ctx->stream_res.audio->enabled == true)
 
                pipe_ctx->stream_res.audio->funcs->az_enable(pipe_ctx->stream_res.audio);
 
-               if (num_audio >= 1 && pp_smu != NULL)
+               if (num_audio >= 1 && clk_mgr->funcs->enable_pme_wa)
                        /*this is the first audio. apply the PME w/a in order to wake AZ from D3*/
-                       set_pme_wa_enable_by_version(core_dc);
+                       clk_mgr->funcs->enable_pme_wa(clk_mgr);
                /* un-mute audio */
                /* TODO: audio should be per stream rather than per link */
                pipe_ctx->stream_res.stream_enc->funcs->audio_mute_control(
 {
        struct dc *dc = pipe_ctx->stream->ctx->dc;
        struct pp_smu_funcs *pp_smu = NULL;
+       struct clk_mgr *clk_mgr = dc->clk_mgr;
 
        if (pipe_ctx->stream_res.audio && pipe_ctx->stream_res.audio->enabled == false)
                return;
                        update_audio_usage(&dc->current_state->res_ctx, dc->res_pool, pipe_ctx->stream_res.audio, false);
                        pipe_ctx->stream_res.audio = NULL;
                }
-               if (pp_smu != NULL)
+               if (clk_mgr->funcs->enable_pme_wa)
                        /*this is the first audio. apply the PME w/a in order to wake AZ from D3*/
-                       set_pme_wa_enable_by_version(dc);
+                       clk_mgr->funcs->enable_pme_wa(clk_mgr);
 
                /* TODO: notify audio driver for if audio modes list changed
                 * add audio mode list change flag */