diff --git a/msm/msm_atomic.c b/msm/msm_atomic.c
index a6e656bbe9..b7bf070977 100644
--- a/msm/msm_atomic.c
+++ b/msm/msm_atomic.c
@@ -419,7 +419,7 @@ static void msm_atomic_helper_commit_modeset_enables(struct drm_device *dev,
 					 crtc->base.id);
 
 			if (funcs->atomic_enable)
-				funcs->atomic_enable(crtc, old_crtc_state);
+				funcs->atomic_enable(crtc, old_state);
 			else
 				funcs->commit(crtc);
 		}
diff --git a/msm/sde/sde_connector.c b/msm/sde/sde_connector.c
index c27cfdce42..2fbb82445c 100644
--- a/msm/sde/sde_connector.c
+++ b/msm/sde/sde_connector.c
@@ -2620,10 +2620,12 @@ sde_connector_best_encoder(struct drm_connector *connector)
 
 static struct drm_encoder *
 sde_connector_atomic_best_encoder(struct drm_connector *connector,
-		struct drm_connector_state *connector_state)
+		struct drm_atomic_state *state)
 {
 	struct sde_connector *c_conn;
 	struct drm_encoder *encoder = NULL;
+	struct drm_connector_state *connector_state =
+			drm_atomic_get_new_connector_state(state, connector);
 
 	if (!connector) {
 		SDE_ERROR("invalid connector\n");
diff --git a/msm/sde/sde_crtc.c b/msm/sde/sde_crtc.c
index 239b76b3bb..e1b9d96866 100644
--- a/msm/sde/sde_crtc.c
+++ b/msm/sde/sde_crtc.c
@@ -3534,7 +3534,7 @@ static void _sde_crtc_clear_all_blend_stages(struct sde_crtc *sde_crtc)
 }
 
 static void sde_crtc_atomic_begin(struct drm_crtc *crtc,
-		struct drm_crtc_state *old_state)
+		struct drm_atomic_state *state)
 {
 	struct sde_crtc *sde_crtc;
 	struct drm_encoder *encoder;
@@ -3543,6 +3543,7 @@ static void sde_crtc_atomic_begin(struct drm_crtc *crtc,
 	struct sde_splash_display *splash_display;
 	bool cont_splash_enabled = false;
 	size_t i;
+	struct drm_crtc_state *old_state = drm_atomic_get_new_crtc_state(state, crtc);
 
 	if (!crtc) {
 		SDE_ERROR("invalid crtc\n");
@@ -3636,7 +3637,7 @@ end:
 }
 
 static void sde_crtc_atomic_flush(struct drm_crtc *crtc,
-		struct drm_crtc_state *old_crtc_state)
+		struct drm_atomic_state *state)
 {
 	struct drm_encoder *encoder;
 	struct sde_crtc *sde_crtc;
@@ -4602,7 +4603,7 @@ static void sde_crtc_disable(struct drm_crtc *crtc)
 }
 
 static void sde_crtc_enable(struct drm_crtc *crtc,
-		struct drm_crtc_state *old_crtc_state)
+		struct drm_atomic_state *old_state)
 {
 	struct sde_crtc *sde_crtc;
 	struct drm_encoder *encoder;
@@ -5348,7 +5349,7 @@ static int _sde_crtc_check_plane_layout(struct drm_crtc *crtc,
 }
 
 static int sde_crtc_atomic_check(struct drm_crtc *crtc,
-		struct drm_crtc_state *state)
+		struct drm_atomic_state *atomic_state)
 {
 	struct drm_device *dev;
 	struct sde_crtc *sde_crtc;
@@ -5359,6 +5360,7 @@ static int sde_crtc_atomic_check(struct drm_crtc *crtc,
 	struct sde_multirect_plane_states *multirect_plane = NULL;
 	struct drm_connector *conn;
 	struct drm_connector_list_iter conn_iter;
+	struct drm_crtc_state *state = drm_atomic_get_new_crtc_state(atomic_state, crtc);
 
 	if (!crtc) {
 		SDE_ERROR("invalid crtc\n");
diff --git a/msm/sde/sde_plane.c b/msm/sde/sde_plane.c
index eb90373ab7..043e54f5db 100644
--- a/msm/sde/sde_plane.c
+++ b/msm/sde/sde_plane.c
@@ -2704,11 +2704,12 @@ modeset_update:
 }
 
 static int sde_plane_atomic_check(struct drm_plane *plane,
-		struct drm_plane_state *state)
+		struct drm_atomic_state *atomic_state)
 {
 	int ret = 0;
 	struct sde_plane *psde;
 	struct sde_plane_state *pstate;
+	struct drm_plane_state *state = drm_atomic_get_new_plane_state(atomic_state, plane);
 
 	if (!plane || !state) {
 		SDE_ERROR("invalid arg(s), plane %d state %d\n",
@@ -3392,7 +3393,7 @@ static void _sde_plane_atomic_disable(struct drm_plane *plane,
 				multirect_index, SDE_SSPP_MULTIRECT_TIME_MX);
 }
 
-static void sde_plane_atomic_update(struct drm_plane *plane,
+static void _sde_plane_atomic_update(struct drm_plane *plane,
 				struct drm_plane_state *old_state)
 {
 	struct sde_plane *psde;
@@ -3423,6 +3424,14 @@ static void sde_plane_atomic_update(struct drm_plane *plane,
 	}
 }
 
+static void sde_plane_atomic_update(struct drm_plane *plane,
+				struct drm_atomic_state *atomic_state)
+{
+	struct drm_plane_state *old_state = drm_atomic_get_old_plane_state(atomic_state, plane);
+
+	_sde_plane_atomic_update(plane, old_state);
+}
+
 void sde_plane_restore(struct drm_plane *plane)
 {
 	struct sde_plane *psde;
@@ -3444,7 +3453,7 @@ void sde_plane_restore(struct drm_plane *plane)
 	SDE_DEBUG_PLANE(psde, "\n");
 
 	/* last plane state is same as current state */
-	sde_plane_atomic_update(plane, plane->state);
+	_sde_plane_atomic_update(plane, plane->state);
 }
 
 bool sde_plane_is_cache_required(struct drm_plane *plane,