diff --git a/drivers/gpu/drm/drm_atomic.c b/drivers/gpu/drm/drm_atomic.c
index ff5f034cc405..3208ed603726 100644
--- a/drivers/gpu/drm/drm_atomic.c
+++ b/drivers/gpu/drm/drm_atomic.c
@@ -358,16 +358,12 @@ EXPORT_SYMBOL(drm_atomic_get_connector_state);
  * sequence must be restarted. All other errors are fatal.
  */
 int
-drm_atomic_set_crtc_for_plane(struct drm_atomic_state *state,
-			      struct drm_plane *plane, struct drm_crtc *crtc)
+drm_atomic_set_crtc_for_plane(struct drm_plane_state *plane_state,
+			      struct drm_crtc *crtc)
 {
-	struct drm_plane_state *plane_state =
-			drm_atomic_get_plane_state(state, plane);
+	struct drm_plane *plane = plane_state->plane;
 	struct drm_crtc_state *crtc_state;
 
-	if (WARN_ON(IS_ERR(plane_state)))
-		return PTR_ERR(plane_state);
-
 	if (plane_state->crtc) {
 		crtc_state = drm_atomic_get_crtc_state(plane_state->state,
 						       plane_state->crtc);
diff --git a/drivers/gpu/drm/drm_atomic_helper.c b/drivers/gpu/drm/drm_atomic_helper.c
index 8b32caad89e8..f191da0e8a89 100644
--- a/drivers/gpu/drm/drm_atomic_helper.c
+++ b/drivers/gpu/drm/drm_atomic_helper.c
@@ -1277,7 +1277,7 @@ retry:
 		goto fail;
 	}
 
-	ret = drm_atomic_set_crtc_for_plane(state, plane, crtc);
+	ret = drm_atomic_set_crtc_for_plane(plane_state, crtc);
 	if (ret != 0)
 		goto fail;
 	drm_atomic_set_fb_for_plane(plane_state, fb);
@@ -1356,7 +1356,7 @@ retry:
 		goto fail;
 	}
 
-	ret = drm_atomic_set_crtc_for_plane(state, plane, NULL);
+	ret = drm_atomic_set_crtc_for_plane(plane_state, NULL);
 	if (ret != 0)
 		goto fail;
 	drm_atomic_set_fb_for_plane(plane_state, NULL);
@@ -1519,7 +1519,7 @@ retry:
 
 		crtc_state->enable = false;
 
-		ret = drm_atomic_set_crtc_for_plane(state, crtc->primary, NULL);
+		ret = drm_atomic_set_crtc_for_plane(primary_state, NULL);
 		if (ret != 0)
 			goto fail;
 
@@ -1534,7 +1534,7 @@ retry:
 	crtc_state->enable = true;
 	drm_mode_copy(&crtc_state->mode, set->mode);
 
-	ret = drm_atomic_set_crtc_for_plane(state, crtc->primary, crtc);
+	ret = drm_atomic_set_crtc_for_plane(primary_state, crtc);
 	if (ret != 0)
 		goto fail;
 	drm_atomic_set_fb_for_plane(primary_state, set->fb);
@@ -1806,7 +1806,7 @@ retry:
 		goto fail;
 	}
 
-	ret = drm_atomic_set_crtc_for_plane(state, plane, crtc);
+	ret = drm_atomic_set_crtc_for_plane(plane_state, crtc);
 	if (ret != 0)
 		goto fail;
 	drm_atomic_set_fb_for_plane(plane_state, fb);
@@ -1869,6 +1869,9 @@ void drm_atomic_helper_crtc_reset(struct drm_crtc *crtc)
 {
 	kfree(crtc->state);
 	crtc->state = kzalloc(sizeof(*crtc->state), GFP_KERNEL);
+
+	if (crtc->state)
+		crtc->state->crtc = crtc;
 }
 EXPORT_SYMBOL(drm_atomic_helper_crtc_reset);
 
@@ -1928,6 +1931,9 @@ void drm_atomic_helper_plane_reset(struct drm_plane *plane)
 
 	kfree(plane->state);
 	plane->state = kzalloc(sizeof(*plane->state), GFP_KERNEL);
+
+	if (plane->state)
+		plane->state->plane = plane;
 }
 EXPORT_SYMBOL(drm_atomic_helper_plane_reset);
 
@@ -1985,6 +1991,9 @@ void drm_atomic_helper_connector_reset(struct drm_connector *connector)
 {
 	kfree(connector->state);
 	connector->state = kzalloc(sizeof(*connector->state), GFP_KERNEL);
+
+	if (connector->state)
+		connector->state->connector = connector;
 }
 EXPORT_SYMBOL(drm_atomic_helper_connector_reset);
 
diff --git a/drivers/gpu/drm/drm_crtc_helper.c b/drivers/gpu/drm/drm_crtc_helper.c
index d552708409de..b1979e7bdc88 100644
--- a/drivers/gpu/drm/drm_crtc_helper.c
+++ b/drivers/gpu/drm/drm_crtc_helper.c
@@ -946,6 +946,7 @@ int drm_helper_crtc_mode_set(struct drm_crtc *crtc, struct drm_display_mode *mod
 		crtc_state = kzalloc(sizeof(*crtc_state), GFP_KERNEL);
 	if (!crtc_state)
 		return -ENOMEM;
+	crtc_state->crtc = crtc;
 
 	crtc_state->enable = true;
 	crtc_state->planes_changed = true;
@@ -1005,6 +1006,7 @@ int drm_helper_crtc_mode_set_base(struct drm_crtc *crtc, int x, int y,
 		plane_state = kzalloc(sizeof(*plane_state), GFP_KERNEL);
 	if (!plane_state)
 		return -ENOMEM;
+	plane_state->plane = plane;
 
 	plane_state->crtc = crtc;
 	drm_atomic_set_fb_for_plane(plane_state, crtc->primary->fb);
diff --git a/drivers/gpu/drm/drm_plane_helper.c b/drivers/gpu/drm/drm_plane_helper.c
index 18a1ac6ac22f..ae61fb21ea86 100644
--- a/drivers/gpu/drm/drm_plane_helper.c
+++ b/drivers/gpu/drm/drm_plane_helper.c
@@ -517,6 +517,7 @@ int drm_plane_helper_update(struct drm_plane *plane, struct drm_crtc *crtc,
 		plane_state = kzalloc(sizeof(*plane_state), GFP_KERNEL);
 	if (!plane_state)
 		return -ENOMEM;
+	plane_state->plane = plane;
 
 	plane_state->crtc = crtc;
 	drm_atomic_set_fb_for_plane(plane_state, fb);
@@ -563,6 +564,7 @@ int drm_plane_helper_disable(struct drm_plane *plane)
 		plane_state = kzalloc(sizeof(*plane_state), GFP_KERNEL);
 	if (!plane_state)
 		return -ENOMEM;
+	plane_state->plane = plane;
 
 	plane_state->crtc = NULL;
 	drm_atomic_set_fb_for_plane(plane_state, NULL);
diff --git a/drivers/gpu/drm/msm/mdp/mdp5/mdp5_plane.c b/drivers/gpu/drm/msm/mdp/mdp5/mdp5_plane.c
index 26e5fdea6594..fc76f630e5b1 100644
--- a/drivers/gpu/drm/msm/mdp/mdp5/mdp5_plane.c
+++ b/drivers/gpu/drm/msm/mdp/mdp5/mdp5_plane.c
@@ -113,6 +113,7 @@ static void mdp5_plane_reset(struct drm_plane *plane)
 	} else {
 		mdp5_state->zpos = 1 + drm_plane_index(plane);
 	}
+	mdp5_state->base.plane = plane;
 
 	plane->state = &mdp5_state->base;
 }
diff --git a/include/drm/drm_atomic.h b/include/drm/drm_atomic.h
index ad2229574dd9..e224ccfa11ca 100644
--- a/include/drm/drm_atomic.h
+++ b/include/drm/drm_atomic.h
@@ -46,8 +46,8 @@ drm_atomic_get_connector_state(struct drm_atomic_state *state,
 			       struct drm_connector *connector);
 
 int __must_check
-drm_atomic_set_crtc_for_plane(struct drm_atomic_state *state,
-			      struct drm_plane *plane, struct drm_crtc *crtc);
+drm_atomic_set_crtc_for_plane(struct drm_plane_state *plane_state,
+			      struct drm_crtc *crtc);
 void drm_atomic_set_fb_for_plane(struct drm_plane_state *plane_state,
 				 struct drm_framebuffer *fb);
 int __must_check
diff --git a/include/drm/drm_crtc.h b/include/drm/drm_crtc.h
index 5006f799ecca..43cdf7d60713 100644
--- a/include/drm/drm_crtc.h
+++ b/include/drm/drm_crtc.h
@@ -229,6 +229,7 @@ struct drm_atomic_state;
 
 /**
  * struct drm_crtc_state - mutable CRTC state
+ * @crtc: backpointer to the CRTC
  * @enable: whether the CRTC should be enabled, gates all other state
  * @mode_changed: for use by helpers and drivers when computing state updates
  * @plane_mask: bitmask of (1 << drm_plane_index(plane)) of attached planes
@@ -242,6 +243,8 @@ struct drm_atomic_state;
  * @state: backpointer to global drm_atomic_state
  */
 struct drm_crtc_state {
+	struct drm_crtc *crtc;
+
 	bool enable;
 
 	/* computed state bits used by helpers and drivers */
@@ -441,11 +444,14 @@ struct drm_crtc {
 
 /**
  * struct drm_connector_state - mutable connector state
+ * @connector: backpointer to the connector
  * @crtc: CRTC to connect connector to, NULL if disabled
  * @best_encoder: can be used by helpers and drivers to select the encoder
  * @state: backpointer to global drm_atomic_state
  */
 struct drm_connector_state {
+	struct drm_connector *connector;
+
 	struct drm_crtc *crtc;  /* do not write directly, use drm_atomic_set_crtc_for_connector() */
 
 	struct drm_encoder *best_encoder;
@@ -665,6 +671,7 @@ struct drm_connector {
 
 /**
  * struct drm_plane_state - mutable plane state
+ * @plane: backpointer to the plane
  * @crtc: currently bound CRTC, NULL if disabled
  * @fb: currently bound framebuffer
  * @fence: optional fence to wait for before scanning out @fb
@@ -681,6 +688,8 @@ struct drm_connector {
  * @state: backpointer to global drm_atomic_state
  */
 struct drm_plane_state {
+	struct drm_plane *plane;
+
 	struct drm_crtc *crtc;   /* do not write directly, use drm_atomic_set_crtc_for_plane() */
 	struct drm_framebuffer *fb;  /* do not write directly, use drm_atomic_set_fb_for_plane() */
 	struct fence *fence;
