diff --git a/drivers/media/media-device.c b/drivers/media/media-device.c
index 17cd349e485f..a2353ecab687 100644
--- a/drivers/media/media-device.c
+++ b/drivers/media/media-device.c
@@ -589,6 +589,26 @@ int __must_check media_device_register_entity(struct media_device *mdev,
 
 	spin_unlock(&mdev->lock);
 
+	mutex_lock(&mdev->graph_mutex);
+	if (mdev->entity_internal_idx_max
+	    >= mdev->pm_count_walk.ent_enum.idx_max) {
+		struct media_entity_graph new = { .top = 0 };
+
+		/*
+		 * Initialise the new graph walk before cleaning up
+		 * the old one in order not to spoil the graph walk
+		 * object of the media device if graph walk init fails.
+		 */
+		ret = media_entity_graph_walk_init(&new, mdev);
+		if (ret) {
+			mutex_unlock(&mdev->graph_mutex);
+			return ret;
+		}
+		media_entity_graph_walk_cleanup(&mdev->pm_count_walk);
+		mdev->pm_count_walk = new;
+	}
+	mutex_unlock(&mdev->graph_mutex);
+
 	return 0;
 }
 EXPORT_SYMBOL_GPL(media_device_register_entity);
@@ -667,6 +687,7 @@ void media_device_cleanup(struct media_device *mdev)
 {
 	ida_destroy(&mdev->entity_internal_idx);
 	mdev->entity_internal_idx_max = 0;
+	media_entity_graph_walk_cleanup(&mdev->pm_count_walk);
 	mutex_destroy(&mdev->graph_mutex);
 }
 EXPORT_SYMBOL_GPL(media_device_cleanup);
diff --git a/include/media/media-device.h b/include/media/media-device.h
index 32bbba7d8115..df74cfa7da4a 100644
--- a/include/media/media-device.h
+++ b/include/media/media-device.h
@@ -306,6 +306,8 @@ struct media_entity_notify {
  * @entity_notify: List of registered entity_notify callbacks
  * @lock:	Entities list lock
  * @graph_mutex: Entities graph operation lock
+ * @pm_count_walk: Graph walk for power state walk. Access serialised using
+ *		   graph_mutex.
  *
  * @source_priv: Driver Private data for enable/disable source handlers
  * @enable_source: Enable Source Handler function pointer
@@ -373,6 +375,7 @@ struct media_device {
 	spinlock_t lock;
 	/* Serializes graph operations. */
 	struct mutex graph_mutex;
+	struct media_entity_graph pm_count_walk;
 
 	void *source_priv;
 	int (*enable_source)(struct media_entity *entity,