return display_count;
 }
 
-static void dcn31_disable_otg_wa(struct clk_mgr *clk_mgr_base, struct dc_state *context, bool disable)
+static void dcn31_disable_otg_wa(struct clk_mgr *clk_mgr_base, bool disable)
 {
        struct dc *dc = clk_mgr_base->ctx->dc;
        int i;
 
        for (i = 0; i < dc->res_pool->pipe_count; ++i) {
-               struct pipe_ctx *pipe = &context->res_ctx.pipe_ctx[i];
+               struct pipe_ctx *pipe = &dc->current_state->res_ctx.pipe_ctx[i];
 
                if (pipe->top_pipe || pipe->prev_odm_pipe)
                        continue;
        }
 
        if (should_set_clock(safe_to_lower, new_clocks->dispclk_khz, clk_mgr_base->clks.dispclk_khz)) {
-               dcn31_disable_otg_wa(clk_mgr_base, context, true);
+               dcn31_disable_otg_wa(clk_mgr_base, true);
 
                clk_mgr_base->clks.dispclk_khz = new_clocks->dispclk_khz;
                dcn31_smu_set_dispclk(clk_mgr, clk_mgr_base->clks.dispclk_khz);
-               dcn31_disable_otg_wa(clk_mgr_base, context, false);
+               dcn31_disable_otg_wa(clk_mgr_base, false);
 
                update_dispclk = true;
        }
 
        return display_count;
 }
 
-static void dcn315_disable_otg_wa(struct clk_mgr *clk_mgr_base, struct dc_state *context, bool disable)
+static void dcn315_disable_otg_wa(struct clk_mgr *clk_mgr_base, bool disable)
 {
        struct dc *dc = clk_mgr_base->ctx->dc;
        int i;
 
        for (i = 0; i < dc->res_pool->pipe_count; ++i) {
-               struct pipe_ctx *pipe = &context->res_ctx.pipe_ctx[i];
+               struct pipe_ctx *pipe = &dc->current_state->res_ctx.pipe_ctx[i];
 
                if (pipe->top_pipe || pipe->prev_odm_pipe)
                        continue;
        }
 
        if (should_set_clock(safe_to_lower, new_clocks->dispclk_khz, clk_mgr_base->clks.dispclk_khz)) {
-               dcn315_disable_otg_wa(clk_mgr_base, context, true);
+               dcn315_disable_otg_wa(clk_mgr_base, true);
 
                clk_mgr_base->clks.dispclk_khz = new_clocks->dispclk_khz;
                dcn315_smu_set_dispclk(clk_mgr, clk_mgr_base->clks.dispclk_khz);
-               dcn315_disable_otg_wa(clk_mgr_base, context, false);
+               dcn315_disable_otg_wa(clk_mgr_base, false);
 
                update_dispclk = true;
        }
 
        return display_count;
 }
 
-static void dcn316_disable_otg_wa(struct clk_mgr *clk_mgr_base, struct dc_state *context, bool disable)
+static void dcn316_disable_otg_wa(struct clk_mgr *clk_mgr_base, bool disable)
 {
        struct dc *dc = clk_mgr_base->ctx->dc;
        int i;
 
        for (i = 0; i < dc->res_pool->pipe_count; ++i) {
-               struct pipe_ctx *pipe = &context->res_ctx.pipe_ctx[i];
+               struct pipe_ctx *pipe = &dc->current_state->res_ctx.pipe_ctx[i];
 
                if (pipe->top_pipe || pipe->prev_odm_pipe)
                        continue;
        }
 
        if (should_set_clock(safe_to_lower, new_clocks->dispclk_khz, clk_mgr_base->clks.dispclk_khz)) {
-               dcn316_disable_otg_wa(clk_mgr_base, context, true);
+               dcn316_disable_otg_wa(clk_mgr_base, true);
 
                clk_mgr_base->clks.dispclk_khz = new_clocks->dispclk_khz;
                dcn316_smu_set_dispclk(clk_mgr, clk_mgr_base->clks.dispclk_khz);
-               dcn316_disable_otg_wa(clk_mgr_base, context, false);
+               dcn316_disable_otg_wa(clk_mgr_base, false);
 
                update_dispclk = true;
        }