diff --git a/drivers/gpu/drm/tegra/dc.c b/drivers/gpu/drm/tegra/dc.c
index dfa787e383cc7ba31519bc3d85ca1f063b815011..464495b46db2c92420091c5f285ef4753e772ef0 100644
--- a/drivers/gpu/drm/tegra/dc.c
+++ b/drivers/gpu/drm/tegra/dc.c
@@ -1736,31 +1736,6 @@ static void tegra_crtc_atomic_enable(struct drm_crtc *crtc,
 	drm_crtc_vblank_on(crtc);
 }
 
-static int tegra_crtc_atomic_check(struct drm_crtc *crtc,
-				   struct drm_crtc_state *state)
-{
-	struct tegra_atomic_state *s = to_tegra_atomic_state(state->state);
-	struct tegra_dc_state *tegra = to_dc_state(state);
-
-	/*
-	 * The display hub display clock needs to be fed by the display clock
-	 * with the highest frequency to ensure proper functioning of all the
-	 * displays.
-	 *
-	 * Note that this isn't used before Tegra186, but it doesn't hurt and
-	 * conditionalizing it would make the code less clean.
-	 */
-	if (state->active) {
-		if (!s->clk_disp || tegra->pclk > s->rate) {
-			s->dc = to_tegra_dc(crtc);
-			s->clk_disp = s->dc->clk;
-			s->rate = tegra->pclk;
-		}
-	}
-
-	return 0;
-}
-
 static void tegra_crtc_atomic_begin(struct drm_crtc *crtc,
 				    struct drm_crtc_state *old_crtc_state)
 {
@@ -1797,7 +1772,6 @@ static void tegra_crtc_atomic_flush(struct drm_crtc *crtc,
 }
 
 static const struct drm_crtc_helper_funcs tegra_crtc_helper_funcs = {
-	.atomic_check = tegra_crtc_atomic_check,
 	.atomic_begin = tegra_crtc_atomic_begin,
 	.atomic_flush = tegra_crtc_atomic_flush,
 	.atomic_enable = tegra_crtc_atomic_enable,
diff --git a/drivers/gpu/drm/tegra/drm.c b/drivers/gpu/drm/tegra/drm.c
index d50bddb2e4474e456b8105d6b90b47488ff21263..e20e013151f048eb1f7c9c246fbba027dd015f82 100644
--- a/drivers/gpu/drm/tegra/drm.c
+++ b/drivers/gpu/drm/tegra/drm.c
@@ -42,6 +42,10 @@ static int tegra_atomic_check(struct drm_device *drm,
 	if (err < 0)
 		return err;
 
+	err = tegra_display_hub_atomic_check(drm, state);
+	if (err < 0)
+		return err;
+
 	err = drm_atomic_normalize_zpos(drm, state);
 	if (err < 0)
 		return err;
@@ -56,35 +60,6 @@ static int tegra_atomic_check(struct drm_device *drm,
 	return 0;
 }
 
-static struct drm_atomic_state *
-tegra_atomic_state_alloc(struct drm_device *drm)
-{
-	struct tegra_atomic_state *state = kzalloc(sizeof(*state), GFP_KERNEL);
-
-	if (!state || drm_atomic_state_init(drm, &state->base) < 0) {
-		kfree(state);
-		return NULL;
-	}
-
-	return &state->base;
-}
-
-static void tegra_atomic_state_clear(struct drm_atomic_state *state)
-{
-	struct tegra_atomic_state *tegra = to_tegra_atomic_state(state);
-
-	drm_atomic_state_default_clear(state);
-	tegra->clk_disp = NULL;
-	tegra->dc = NULL;
-	tegra->rate = 0;
-}
-
-static void tegra_atomic_state_free(struct drm_atomic_state *state)
-{
-	drm_atomic_state_default_release(state);
-	kfree(state);
-}
-
 static const struct drm_mode_config_funcs tegra_drm_mode_config_funcs = {
 	.fb_create = tegra_fb_create,
 #ifdef CONFIG_DRM_FBDEV_EMULATION
@@ -92,9 +67,6 @@ static const struct drm_mode_config_funcs tegra_drm_mode_config_funcs = {
 #endif
 	.atomic_check = tegra_atomic_check,
 	.atomic_commit = drm_atomic_helper_commit,
-	.atomic_state_alloc = tegra_atomic_state_alloc,
-	.atomic_state_clear = tegra_atomic_state_clear,
-	.atomic_state_free = tegra_atomic_state_free,
 };
 
 static void tegra_atomic_commit_tail(struct drm_atomic_state *old_state)
diff --git a/drivers/gpu/drm/tegra/drm.h b/drivers/gpu/drm/tegra/drm.h
index 73b661ce7086592af497b39bdaf94fe806b8860f..4f41aaec85301c486558028b9a2a509e7d73fa8a 100644
--- a/drivers/gpu/drm/tegra/drm.h
+++ b/drivers/gpu/drm/tegra/drm.h
@@ -42,20 +42,6 @@ struct tegra_fbdev {
 };
 #endif
 
-struct tegra_atomic_state {
-	struct drm_atomic_state base;
-
-	struct clk *clk_disp;
-	struct tegra_dc *dc;
-	unsigned long rate;
-};
-
-static inline struct tegra_atomic_state *
-to_tegra_atomic_state(struct drm_atomic_state *state)
-{
-	return container_of(state, struct tegra_atomic_state, base);
-}
-
 struct tegra_drm {
 	struct drm_device *drm;
 
diff --git a/drivers/gpu/drm/tegra/hub.c b/drivers/gpu/drm/tegra/hub.c
index 094324daa917a37bb2888707036c45271ddc7dd7..9a3f23d4780fd4eac622572ea2016812d33ec530 100644
--- a/drivers/gpu/drm/tegra/hub.c
+++ b/drivers/gpu/drm/tegra/hub.c
@@ -573,6 +573,89 @@ struct drm_plane *tegra_shared_plane_create(struct drm_device *drm,
 	return p;
 }
 
+static struct drm_private_state *
+tegra_display_hub_duplicate_state(struct drm_private_obj *obj)
+{
+	struct tegra_display_hub_state *state;
+
+	state = kmemdup(obj->state, sizeof(*state), GFP_KERNEL);
+	if (!state)
+		return NULL;
+
+	__drm_atomic_helper_private_obj_duplicate_state(obj, &state->base);
+
+	return &state->base;
+}
+
+static void tegra_display_hub_destroy_state(struct drm_private_obj *obj,
+					    struct drm_private_state *state)
+{
+	struct tegra_display_hub_state *hub_state =
+		to_tegra_display_hub_state(state);
+
+	kfree(hub_state);
+}
+
+static const struct drm_private_state_funcs tegra_display_hub_state_funcs = {
+	.atomic_duplicate_state = tegra_display_hub_duplicate_state,
+	.atomic_destroy_state = tegra_display_hub_destroy_state,
+};
+
+static struct tegra_display_hub_state *
+tegra_display_hub_get_state(struct tegra_display_hub *hub,
+			    struct drm_atomic_state *state)
+{
+	struct drm_device *drm = dev_get_drvdata(hub->client.parent);
+	struct drm_private_state *priv;
+
+	WARN_ON(!drm_modeset_is_locked(&drm->mode_config.connection_mutex));
+
+	priv = drm_atomic_get_private_obj_state(state, &hub->base);
+	if (IS_ERR(priv))
+		return ERR_CAST(priv);
+
+	return to_tegra_display_hub_state(priv);
+}
+
+int tegra_display_hub_atomic_check(struct drm_device *drm,
+				   struct drm_atomic_state *state)
+{
+	struct tegra_drm *tegra = drm->dev_private;
+	struct tegra_display_hub_state *hub_state;
+	struct drm_crtc_state *old, *new;
+	struct drm_crtc *crtc;
+	unsigned int i;
+
+	if (!tegra->hub)
+		return 0;
+
+	hub_state = tegra_display_hub_get_state(tegra->hub, state);
+	if (IS_ERR(hub_state))
+		return PTR_ERR(hub_state);
+
+	/*
+	 * The display hub display clock needs to be fed by the display clock
+	 * with the highest frequency to ensure proper functioning of all the
+	 * displays.
+	 *
+	 * Note that this isn't used before Tegra186, but it doesn't hurt and
+	 * conditionalizing it would make the code less clean.
+	 */
+	for_each_oldnew_crtc_in_state(state, crtc, old, new, i) {
+		struct tegra_dc_state *dc = to_dc_state(new);
+
+		if (new->active) {
+			if (!hub_state->clk || dc->pclk > hub_state->rate) {
+				hub_state->dc = to_tegra_dc(dc->base.crtc);
+				hub_state->clk = hub_state->dc->clk;
+				hub_state->rate = dc->pclk;
+			}
+		}
+	}
+
+	return 0;
+}
+
 static void tegra_display_hub_update(struct tegra_dc *dc)
 {
 	u32 value;
@@ -598,26 +681,28 @@ static void tegra_display_hub_update(struct tegra_dc *dc)
 void tegra_display_hub_atomic_commit(struct drm_device *drm,
 				     struct drm_atomic_state *state)
 {
-	struct tegra_atomic_state *s = to_tegra_atomic_state(state);
 	struct tegra_drm *tegra = drm->dev_private;
 	struct tegra_display_hub *hub = tegra->hub;
+	struct tegra_display_hub_state *hub_state;
 	struct device *dev = hub->client.dev;
 	int err;
 
-	if (s->clk_disp) {
-		err = clk_set_rate(s->clk_disp, s->rate);
+	hub_state = tegra_display_hub_get_state(hub, state);
+
+	if (hub_state->clk) {
+		err = clk_set_rate(hub_state->clk, hub_state->rate);
 		if (err < 0)
 			dev_err(dev, "failed to set rate of %pC to %lu Hz\n",
-				s->clk_disp, s->rate);
+				hub_state->clk, hub_state->rate);
 
-		err = clk_set_parent(hub->clk_disp, s->clk_disp);
+		err = clk_set_parent(hub->clk_disp, hub_state->clk);
 		if (err < 0)
 			dev_err(dev, "failed to set parent of %pC to %pC: %d\n",
-				hub->clk_disp, s->clk_disp, err);
+				hub->clk_disp, hub_state->clk, err);
 	}
 
-	if (s->dc)
-		tegra_display_hub_update(s->dc);
+	if (hub_state->dc)
+		tegra_display_hub_update(hub_state->dc);
 }
 
 static int tegra_display_hub_init(struct host1x_client *client)
@@ -625,6 +710,14 @@ static int tegra_display_hub_init(struct host1x_client *client)
 	struct tegra_display_hub *hub = to_tegra_display_hub(client);
 	struct drm_device *drm = dev_get_drvdata(client->parent);
 	struct tegra_drm *tegra = drm->dev_private;
+	struct tegra_display_hub_state *state;
+
+	state = kzalloc(sizeof(*state), GFP_KERNEL);
+	if (!state)
+		return -ENOMEM;
+
+	drm_atomic_private_obj_init(&hub->base, &state->base,
+				    &tegra_display_hub_state_funcs);
 
 	tegra->hub = hub;
 
@@ -636,6 +729,7 @@ static int tegra_display_hub_exit(struct host1x_client *client)
 	struct drm_device *drm = dev_get_drvdata(client->parent);
 	struct tegra_drm *tegra = drm->dev_private;
 
+	drm_atomic_private_obj_fini(&tegra->hub->base);
 	tegra->hub = NULL;
 
 	return 0;
diff --git a/drivers/gpu/drm/tegra/hub.h b/drivers/gpu/drm/tegra/hub.h
index 890a47cd05c310aa507d2a95a363dfaec73471da..85b8bf41a395ac748fdeb5c48db325123fd92a29 100644
--- a/drivers/gpu/drm/tegra/hub.h
+++ b/drivers/gpu/drm/tegra/hub.h
@@ -41,6 +41,7 @@ struct tegra_display_hub_soc {
 };
 
 struct tegra_display_hub {
+	struct drm_private_obj base;
 	struct host1x_client client;
 	struct clk *clk_disp;
 	struct clk *clk_dsc;
@@ -57,6 +58,20 @@ to_tegra_display_hub(struct host1x_client *client)
 	return container_of(client, struct tegra_display_hub, client);
 }
 
+struct tegra_display_hub_state {
+	struct drm_private_state base;
+
+	struct tegra_dc *dc;
+	unsigned long rate;
+	struct clk *clk;
+};
+
+static inline struct tegra_display_hub_state *
+to_tegra_display_hub_state(struct drm_private_state *priv)
+{
+	return container_of(priv, struct tegra_display_hub_state, base);
+}
+
 struct tegra_dc;
 struct tegra_plane;
 
@@ -68,6 +83,8 @@ struct drm_plane *tegra_shared_plane_create(struct drm_device *drm,
 					    unsigned int wgrp,
 					    unsigned int index);
 
+int tegra_display_hub_atomic_check(struct drm_device *drm,
+				   struct drm_atomic_state *state);
 void tegra_display_hub_atomic_commit(struct drm_device *drm,
 				     struct drm_atomic_state *state);