dpu_core_irq_uninstall(dpu_kms);
 }
 
-static void dpu_kms_mdp_snapshot(struct msm_kms *kms)
+static void dpu_kms_mdp_snapshot(struct msm_disp_state *disp_state, struct msm_kms *kms)
 {
        int i;
        struct dpu_kms *dpu_kms;
        struct dpu_mdss_cfg *cat;
        struct dpu_hw_mdp *top;
-       struct msm_disp_state *disp_state;
-
-       disp_state = kms->disp_state;
 
        dpu_kms = to_dpu_kms(kms);
 
 
  */
 void msm_disp_snapshot_state(struct drm_device *drm_dev);
 
-/**
- * msm_disp_state_get - get the handle to msm_disp_state struct from the drm device
- * @drm:           handle to drm device
-
- * Returns:    handle to the msm_disp_state struct
- */
-struct msm_disp_state *msm_disp_state_get(struct drm_device *drm);
-
 /**
  * msm_disp_state_print - print out the current dpu state
  * @disp_state:            handle to drm device
 
        }
 }
 
-struct msm_disp_state *msm_disp_state_get(struct drm_device *drm)
-{
-       struct msm_drm_private *priv;
-       struct msm_kms *kms;
-
-       priv = drm->dev_private;
-       kms = priv->kms;
-
-       return kms->disp_state;
-}
-
 void msm_disp_state_print(struct msm_disp_state *state, struct drm_printer *p)
 {
        struct msm_disp_state_block *block, *tmp;
        kms = priv->kms;
 
        if (priv->dp)
-               msm_dp_snapshot(priv->dp);
+               msm_dp_snapshot(disp_state, priv->dp);
 
        for (i = 0; i < ARRAY_SIZE(priv->dsi); i++) {
                if (!priv->dsi[i])
                        continue;
 
-               msm_dsi_snapshot(priv->dsi[i]);
+               msm_dsi_snapshot(disp_state, priv->dsi[i]);
        }
 
        if (kms->funcs->snapshot)
-               kms->funcs->snapshot(kms);
+               kms->funcs->snapshot(disp_state, kms);
 
        msm_disp_capture_atomic_state(disp_state);
 }
 
                dp_display->link->test_video.test_bit_depth);
 }
 
-void msm_dp_snapshot(struct msm_dp *dp)
+void msm_dp_snapshot(struct msm_disp_state *disp_state, struct msm_dp *dp)
 {
        struct dp_display_private *dp_display;
        struct drm_device *drm;
-       struct msm_disp_state *disp_state;
 
        dp_display = container_of(dp, struct dp_display_private, dp_display);
        drm = dp->drm_dev;
-       disp_state = msm_disp_state_get(drm);
 
        /*
         * if we are reading registers we need the link clocks to be on
 
        return ret;
 }
 
-void msm_dsi_snapshot(struct msm_dsi *msm_dsi)
+void msm_dsi_snapshot(struct msm_disp_state *disp_state, struct msm_dsi *msm_dsi)
 {
-       msm_dsi_host_snapshot(msm_dsi->host);
+       msm_dsi_host_snapshot(disp_state, msm_dsi->host);
 }
 
 
        return msm_dsi->panel || msm_dsi->external_bridge;
 }
 
-void msm_dsi_snapshot(struct msm_dsi *msm_dsi);
+void msm_dsi_snapshot(struct msm_disp_state *disp_state, struct msm_dsi *msm_dsi);
 
 struct drm_encoder *msm_dsi_get_encoder(struct msm_dsi *msm_dsi);
 
 int dsi_clk_init_6g_v2(struct msm_dsi_host *msm_host);
 int dsi_calc_clk_rate_v2(struct msm_dsi_host *msm_host, bool is_dual_dsi);
 int dsi_calc_clk_rate_6g(struct msm_dsi_host *msm_host, bool is_dual_dsi);
-void msm_dsi_host_snapshot(struct mipi_dsi_host *host);
+void msm_dsi_host_snapshot(struct msm_disp_state *disp_state, struct mipi_dsi_host *host);
 /* dsi phy */
 struct msm_dsi_phy;
 struct msm_dsi_phy_shared_timings {
 
        return of_drm_find_bridge(msm_host->device_node);
 }
 
-void msm_dsi_host_snapshot(struct mipi_dsi_host *host)
+void msm_dsi_host_snapshot(struct msm_disp_state *disp_state, struct mipi_dsi_host *host)
 {
        struct msm_dsi_host *msm_host = to_msm_dsi_host(host);
-       struct drm_device *dev = msm_host->dev;
-       struct msm_disp_state *disp_state;
-
-       disp_state = msm_disp_state_get(dev);
 
        pm_runtime_get_sync(&msm_host->pdev->dev);
 
 
                                struct drm_display_mode *mode,
                                struct drm_display_mode *adjusted_mode);
 void msm_dp_irq_postinstall(struct msm_dp *dp_display);
-void msm_dp_snapshot(struct msm_dp *dp_display);
+struct msm_disp_state;
+void msm_dp_snapshot(struct msm_disp_state *disp_state, struct msm_dp *dp_display);
 
 void msm_dp_debugfs_init(struct msm_dp *dp_display, struct drm_minor *minor);
 
 
        void (*destroy)(struct msm_kms *kms);
 
        /* snapshot: */
-       void (*snapshot)(struct msm_kms *kms);
+       void (*snapshot)(struct msm_disp_state *disp_state, struct msm_kms *kms);
 
 #ifdef CONFIG_DEBUG_FS
        /* debugfs: */