return pipe_cnt;
 }
 
-void dcn31_update_soc_for_wm_a(struct dc *dc, struct dc_state *context)
-{
-       if (dc->clk_mgr->bw_params->wm_table.entries[WM_A].valid) {
-               context->bw_ctx.dml.soc.dram_clock_change_latency_us = dc->clk_mgr->bw_params->wm_table.entries[WM_A].pstate_latency_us;
-               context->bw_ctx.dml.soc.sr_enter_plus_exit_time_us = dc->clk_mgr->bw_params->wm_table.entries[WM_A].sr_enter_plus_exit_time_us;
-               context->bw_ctx.dml.soc.sr_exit_time_us = dc->clk_mgr->bw_params->wm_table.entries[WM_A].sr_exit_time_us;
-       }
-}
-
 void dcn31_calculate_wm_and_dlg(
                struct dc *dc, struct dc_state *context,
                display_e2e_pipe_params_st *pipes,
 
                          struct dc_state *context,
                          display_e2e_pipe_params_st *pipes,
                          int pipe_cnt);
-void dcn31_update_soc_for_wm_a(struct dc *dc, struct dc_state *context);
 
 struct resource_pool *dcn31_create_resource_pool(
                const struct dc_init_data *init_data,
 
 #include "virtual/virtual_stream_encoder.h"
 #include "dce110/dce110_resource.h"
 #include "dml/display_mode_vba.h"
+#include "dml/dcn31/dcn31_fpu.h"
 #include "dcn314/dcn314_dccg.h"
 #include "dcn10/dcn10_resource.h"
 #include "dcn31/dcn31_panel_cntl.h"
 
        .urgent_latency_adjustment_fabric_clock_reference_mhz = 0,
 };
 
+void dcn31_update_soc_for_wm_a(struct dc *dc, struct dc_state *context)
+{
+       if (dc->clk_mgr->bw_params->wm_table.entries[WM_A].valid) {
+               context->bw_ctx.dml.soc.dram_clock_change_latency_us = dc->clk_mgr->bw_params->wm_table.entries[WM_A].pstate_latency_us;
+               context->bw_ctx.dml.soc.sr_enter_plus_exit_time_us = dc->clk_mgr->bw_params->wm_table.entries[WM_A].sr_enter_plus_exit_time_us;
+               context->bw_ctx.dml.soc.sr_exit_time_us = dc->clk_mgr->bw_params->wm_table.entries[WM_A].sr_exit_time_us;
+       }
+}
+
 void dcn31_calculate_wm_and_dlg_fp(
                struct dc *dc, struct dc_state *context,
                display_e2e_pipe_params_st *pipes,
 
 #define DCN3_15_MIN_COMPBUF_SIZE_KB 128
 #define DCN3_16_DEFAULT_DET_SIZE 192
 
+void dcn31_update_soc_for_wm_a(struct dc *dc, struct dc_state *context);
+
 void dcn31_calculate_wm_and_dlg_fp(
                struct dc *dc, struct dc_state *context,
                display_e2e_pipe_params_st *pipes,