diff --git a/drivers/gpu/drm/drm_atomic.c b/drivers/gpu/drm/drm_atomic.c
index ba49b5ca822f..ff5f034cc405 100644
--- a/drivers/gpu/drm/drm_atomic.c
+++ b/drivers/gpu/drm/drm_atomic.c
@@ -344,7 +344,8 @@ EXPORT_SYMBOL(drm_atomic_get_connector_state);
 
 /**
  * drm_atomic_set_crtc_for_plane - set crtc for plane
- * @plane_state: atomic state object for the plane
+ * @state: the incoming atomic state
+ * @plane: the plane whose incoming state to update
  * @crtc: crtc to use for the plane
  *
  * Changing the assigned crtc for a plane requires us to grab the lock and state
@@ -357,20 +358,35 @@ 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_plane_state *plane_state,
-			      struct drm_crtc *crtc)
+drm_atomic_set_crtc_for_plane(struct drm_atomic_state *state,
+			      struct drm_plane *plane, struct drm_crtc *crtc)
 {
+	struct drm_plane_state *plane_state =
+			drm_atomic_get_plane_state(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);
+		if (WARN_ON(IS_ERR(crtc_state)))
+			return PTR_ERR(crtc_state);
+
+		crtc_state->plane_mask &= ~(1 << drm_plane_index(plane));
+	}
+
+	plane_state->crtc = crtc;
+
 	if (crtc) {
 		crtc_state = drm_atomic_get_crtc_state(plane_state->state,
 						       crtc);
 		if (IS_ERR(crtc_state))
 			return PTR_ERR(crtc_state);
+		crtc_state->plane_mask |= (1 << drm_plane_index(plane));
 	}
 
-	plane_state->crtc = crtc;
-
 	if (crtc)
 		DRM_DEBUG_KMS("Link plane state %p to [CRTC:%d]\n",
 			      plane_state, crtc->base.id);
diff --git a/drivers/gpu/drm/drm_atomic_helper.c b/drivers/gpu/drm/drm_atomic_helper.c
index 2fa0840694d0..2ee509c92034 100644
--- a/drivers/gpu/drm/drm_atomic_helper.c
+++ b/drivers/gpu/drm/drm_atomic_helper.c
@@ -1222,7 +1222,7 @@ retry:
 		goto fail;
 	}
 
-	ret = drm_atomic_set_crtc_for_plane(plane_state, crtc);
+	ret = drm_atomic_set_crtc_for_plane(state, plane, crtc);
 	if (ret != 0)
 		goto fail;
 	drm_atomic_set_fb_for_plane(plane_state, fb);
@@ -1301,7 +1301,7 @@ retry:
 		goto fail;
 	}
 
-	ret = drm_atomic_set_crtc_for_plane(plane_state, NULL);
+	ret = drm_atomic_set_crtc_for_plane(state, plane, NULL);
 	if (ret != 0)
 		goto fail;
 	drm_atomic_set_fb_for_plane(plane_state, NULL);
@@ -1472,7 +1472,7 @@ retry:
 		goto fail;
 	}
 
-	ret = drm_atomic_set_crtc_for_plane(primary_state, crtc);
+	ret = drm_atomic_set_crtc_for_plane(state, crtc->primary, crtc);
 	if (ret != 0)
 		goto fail;
 	drm_atomic_set_fb_for_plane(primary_state, set->fb);
@@ -1744,7 +1744,7 @@ retry:
 		goto fail;
 	}
 
-	ret = drm_atomic_set_crtc_for_plane(plane_state, crtc);
+	ret = drm_atomic_set_crtc_for_plane(state, plane, crtc);
 	if (ret != 0)
 		goto fail;
 	drm_atomic_set_fb_for_plane(plane_state, fb);
diff --git a/include/drm/drm_atomic.h b/include/drm/drm_atomic.h
index e224ccfa11ca..ad2229574dd9 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_plane_state *plane_state,
-			      struct drm_crtc *crtc);
+drm_atomic_set_crtc_for_plane(struct drm_atomic_state *state,
+			      struct drm_plane *plane, 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 b459e8fbbc25..4cf6905b57f5 100644
--- a/include/drm/drm_crtc.h
+++ b/include/drm/drm_crtc.h
@@ -231,6 +231,7 @@ struct drm_atomic_state;
  * struct drm_crtc_state - mutable CRTC state
  * @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
  * @last_vblank_count: for helpers and drivers to capture the vblank of the
  * 	update to ensure framebuffer cleanup isn't done too early
  * @planes_changed: for use by helpers and drivers when computing state updates
@@ -247,6 +248,13 @@ struct drm_crtc_state {
 	bool planes_changed : 1;
 	bool mode_changed : 1;
 
+	/* attached planes bitmask:
+	 * WARNING: transitional helpers do not maintain plane_mask so
+	 * drivers not converted over to atomic helpers should not rely
+	 * on plane_mask being accurate!
+	 */
+	u32 plane_mask;
+
 	/* last_vblank_count: for vblank waits before cleanup */
 	u32 last_vblank_count;
 
@@ -438,7 +446,7 @@ struct drm_crtc {
  * @state: backpointer to global drm_atomic_state
  */
 struct drm_connector_state {
-	struct drm_crtc *crtc;
+	struct drm_crtc *crtc;  /* do not write directly, use drm_atomic_set_crtc_for_connector() */
 
 	struct drm_encoder *best_encoder;
 
@@ -673,8 +681,8 @@ struct drm_connector {
  * @state: backpointer to global drm_atomic_state
  */
 struct drm_plane_state {
-	struct drm_crtc *crtc;
-	struct drm_framebuffer *fb;
+	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;
 
 	/* Signed dest location allows it to be partially off screen */