@@ -29,6 +29,8 @@
 #include <dt-bindings/power/rk3399-power.h>
 #include <dt-bindings/power/rk3568-power.h>
 
+struct rockchip_pmu;
+
 struct rockchip_domain_info {
 	const char *name;
 	int pwr_mask;

@@ -39,6 +41,10 @@ struct rockchip_domain_info {
 	bool active_wakeup;
 	int pwr_w_mask;
 	int req_w_mask;
+
+	struct rockchip_pmu *pmu;
+	struct generic_pm_domain *parent_domain;
+	int id;
 };
 
 struct rockchip_pmu_info {

@@ -81,6 +87,7 @@ struct rockchip_pmu {
 	struct regmap *regmap;
 	const struct rockchip_pmu_info *info;
 	struct mutex mutex; /* mutex lock for pmu */
+	atomic_t missing;
 	struct genpd_onecell_data genpd_data;
 	struct generic_pm_domain *domains[];
 };

@@ -387,12 +394,11 @@ static void rockchip_pd_detach_dev(struct generic_pm_domain *genpd,
 }
 
 static int rockchip_pm_add_one_domain(struct rockchip_pmu *pmu,
-				      struct device_node *node)
+				      struct device_node *node,
+				      struct generic_pm_domain *parent_domain)
 {
-	const struct rockchip_domain_info *pd_info;
-	struct rockchip_pm_domain *pd;
-	struct device_node *qos_node;
-	int i, j;
+	struct platform_device *pd_pdev;
+	struct rockchip_domain_info *domain_info;
 	u32 id;
 	int error;
 

@@ -410,28 +416,98 @@ static int rockchip_pm_add_one_domain(struct rockchip_pmu *pmu,
 		return -EINVAL;
 	}
 
-	pd_info = &pmu->info->domain_info[id];
-	if (!pd_info) {
-		dev_err(pmu->dev, "%pOFn: undefined domain id %d\n",
-			node, id);
-		return -EINVAL;
+	pd_pdev = platform_device_alloc("rk-power-domain", id);
+	if (!pd_pdev) {
+		dev_err(pmu->dev, "Failed to allocate platform device\n");
+		return -ENOMEM;
 	}
 
-	pd = devm_kzalloc(pmu->dev, sizeof(*pd), GFP_KERNEL);
+	error = platform_device_add_data(pd_pdev,
+				         &pmu->info->domain_info[id],
+				         sizeof(pmu->info->domain_info[id]));
+	if (error)
+		goto err_put;
+
+	domain_info = pd_pdev->dev.platform_data;
+	domain_info->parent_domain = parent_domain;
+	domain_info->pmu = pmu;
+	domain_info->id = id;
+
+	pd_pdev->dev.parent = pmu->dev;
+	pd_pdev->dev.of_node = node;
+
+	atomic_inc(&pmu->missing);
+
+	error = platform_device_add(pd_pdev);
+	if (error)
+		goto err_put;
+
+	return 0;
+
+err_put:
+	platform_device_put(pd_pdev);
+
+	return error;
+}
+
+static void rockchip_pm_add_domains(struct rockchip_pmu *pmu,
+				    struct device_node *parent,
+				    struct generic_pm_domain *parent_domain)
+{
+	struct device_node *np;
+	int error;
+
+	/*
+	 * We may only register the genpd provider when we have registered all
+	 * domains that are specified in the device tree. We count the missing
+	 * domains in rockchip_pmu::missing.
+	 * The rockchip pm_domains may have subdomains which means we can be
+	 * called here recursively. Give it one extra count here to prevent
+	 * the counter dropping to zero when we are called recursively.
+	 */
+	atomic_inc(&pmu->missing);
+
+	for_each_child_of_node(parent, np) {
+		error = rockchip_pm_add_one_domain(pmu, np, parent_domain);
+		if (error)
+			dev_err(pmu->dev, "failed to handle node %pOFn: %d\n",
+				np, error);
+	}
+
+	if (!atomic_dec_and_test(&pmu->missing))
+		return;
+
+	error = of_genpd_add_provider_onecell(pmu->dev->of_node, &pmu->genpd_data);
+	if (error)
+		dev_err(pmu->dev, "failed to add provider: %d\n", error);
+}
+
+static int rockchip_domain_probe(struct platform_device *pdev)
+{
+	struct rockchip_domain_info *pd_info = pdev->dev.platform_data;
+	struct rockchip_pm_domain *pd;
+	struct device_node *node = pdev->dev.of_node;
+	struct device_node *qos_node;
+	struct rockchip_pmu *pmu = pd_info->pmu;
+	struct generic_pm_domain *parent_domain;
+	int i, j;
+	int error;
+
+	pd = devm_kzalloc(&pdev->dev, sizeof(*pd), GFP_KERNEL);
 	if (!pd)
 		return -ENOMEM;
 
 	pd->info = pd_info;
-	pd->pmu = pmu;
+	pd->pmu = pd_info->pmu;
 
 	pd->num_clks = of_clk_get_parent_count(node);
 	if (pd->num_clks > 0) {
-		pd->clks = devm_kcalloc(pmu->dev, pd->num_clks,
+		pd->clks = devm_kcalloc(&pdev->dev, pd->num_clks,
 					sizeof(*pd->clks), GFP_KERNEL);
 		if (!pd->clks)
 			return -ENOMEM;
 	} else {
-		dev_dbg(pmu->dev, "%pOFn: doesn't have clocks: %d\n",
+		dev_dbg(&pdev->dev, "%pOFn: doesn't have clocks: %d\n",
 			node, pd->num_clks);
 		pd->num_clks = 0;
 	}

@@ -440,7 +516,7 @@ static int rockchip_pm_add_one_domain(struct rockchip_pmu *pmu,
 		pd->clks[i].clk = of_clk_get(node, i);
 		if (IS_ERR(pd->clks[i].clk)) {
 			error = PTR_ERR(pd->clks[i].clk);
-			dev_err(pmu->dev,
+			dev_err(&pdev->dev,
 				"%pOFn: failed to get clk at index %d: %d\n",
 				node, i, error);
 			return error;

@@ -455,7 +531,7 @@ static int rockchip_pm_add_one_domain(struct rockchip_pmu *pmu,
 						 NULL);
 
 	if (pd->num_qos > 0) {
-		pd->qos_regmap = devm_kcalloc(pmu->dev, pd->num_qos,
+		pd->qos_regmap = devm_kcalloc(&pdev->dev, pd->num_qos,
 					      sizeof(*pd->qos_regmap),
 					      GFP_KERNEL);
 		if (!pd->qos_regmap) {

@@ -464,7 +540,7 @@ static int rockchip_pm_add_one_domain(struct rockchip_pmu *pmu,
 		}
 
 		for (j = 0; j < MAX_QOS_REGS_NUM; j++) {
-			pd->qos_save_regs[j] = devm_kcalloc(pmu->dev,
+			pd->qos_save_regs[j] = devm_kcalloc(&pdev->dev,
 							    pd->num_qos,
 							    sizeof(u32),
 							    GFP_KERNEL);

@@ -490,7 +566,7 @@ static int rockchip_pm_add_one_domain(struct rockchip_pmu *pmu,
 		}
 	}
 
-	error = rockchip_pd_power(pd, true);
+	error = rockchip_pd_power_on(&pd->genpd);
 	if (error) {
 		dev_err(pmu->dev,
 			"failed to power on domain '%pOFn': %d\n",

@@ -507,13 +583,33 @@ static int rockchip_pm_add_one_domain(struct rockchip_pmu *pmu,
 	pd->genpd.attach_dev = rockchip_pd_attach_dev;
 	pd->genpd.detach_dev = rockchip_pd_detach_dev;
 	pd->genpd.flags = GENPD_FLAG_PM_CLK;
-	if (pd_info->active_wakeup)
+	if (pd->info->active_wakeup)
 		pd->genpd.flags |= GENPD_FLAG_ACTIVE_WAKEUP;
 	pm_genpd_init(&pd->genpd, NULL, false);
 
-	pmu->genpd_data.domains[id] = &pd->genpd;
+	parent_domain = pd_info->parent_domain;
+	if (parent_domain) {
+		error = pm_genpd_add_subdomain(parent_domain, &pd->genpd);
+		if (error) {
+			dev_err(pmu->dev, "%s failed to add subdomain %s: %d\n",
+				parent_domain->name, pd->genpd.name, error);
+			goto out_genpd_remove;
+		} else {
+			dev_dbg(pmu->dev, "%s add subdomain: %s\n",
+				parent_domain->name, pd->genpd.name);
+		}
+	}
+
+	pmu->genpd_data.domains[pd->info->id] = &pd->genpd;
+
+	atomic_dec(&pmu->missing);
+
+	rockchip_pm_add_domains(pmu, node, &pd->genpd);
+
 	return 0;
 
+out_genpd_remove:
+	pm_genpd_remove(&pd->genpd);
 err_unprepare_clocks:
 	clk_bulk_unprepare(pd->num_clks, pd->clks);
 err_put_clocks:

@@ -521,46 +617,19 @@ static int rockchip_pm_add_one_domain(struct rockchip_pmu *pmu,
 	return error;
 }
 
-static void rockchip_pm_remove_one_domain(struct rockchip_pm_domain *pd)
+static int rockchip_domain_remove(struct platform_device *pdev)
 {
-	int ret;
-
-	/*
-	 * We're in the error cleanup already, so we only complain,
-	 * but won't emit another error on top of the original one.
-	 */
-	ret = pm_genpd_remove(&pd->genpd);
-	if (ret < 0)
-		dev_err(pd->pmu->dev, "failed to remove domain '%s' : %d - state may be inconsistent\n",
-			pd->genpd.name, ret);
-
-	clk_bulk_unprepare(pd->num_clks, pd->clks);
-	clk_bulk_put(pd->num_clks, pd->clks);
-
-	/* protect the zeroing of pm->num_clks */
-	mutex_lock(&pd->pmu->mutex);
-	pd->num_clks = 0;
-	mutex_unlock(&pd->pmu->mutex);
-
-	/* devm will free our memory */
+	return 0;
 }
 
-static void rockchip_pm_domain_cleanup(struct rockchip_pmu *pmu)
-{
-	struct generic_pm_domain *genpd;
-	struct rockchip_pm_domain *pd;
-	int i;
-
-	for (i = 0; i < pmu->genpd_data.num_domains; i++) {
-		genpd = pmu->genpd_data.domains[i];
-		if (genpd) {
-			pd = to_rockchip_pd(genpd);
-			rockchip_pm_remove_one_domain(pd);
-		}
-	}
-
-	/* devm will free our memory */
-}
+static struct platform_driver rockchip_domain_driver = {
+	.driver = {
+		.name = "rk-power-domain",
+	},
+	.probe    = rockchip_domain_probe,
+	.remove   = rockchip_domain_remove,
+};
+builtin_platform_driver(rockchip_domain_driver)
 
 static void rockchip_configure_pd_cnt(struct rockchip_pmu *pmu,
 				      u32 domain_reg_offset,

@@ -572,71 +641,14 @@ static void rockchip_configure_pd_cnt(struct rockchip_pmu *pmu,
 	regmap_write(pmu->regmap, domain_reg_offset + 4, count);
 }
 
-static int rockchip_pm_add_subdomain(struct rockchip_pmu *pmu,
-				     struct device_node *parent)
-{
-	struct device_node *np;
-	struct generic_pm_domain *child_domain, *parent_domain;
-	int error;
-
-	for_each_child_of_node(parent, np) {
-		u32 idx;
-
-		error = of_property_read_u32(parent, "reg", &idx);
-		if (error) {
-			dev_err(pmu->dev,
-				"%pOFn: failed to retrieve domain id (reg): %d\n",
-				parent, error);
-			goto err_out;
-		}
-		parent_domain = pmu->genpd_data.domains[idx];
-
-		error = rockchip_pm_add_one_domain(pmu, np);
-		if (error) {
-			dev_err(pmu->dev, "failed to handle node %pOFn: %d\n",
-				np, error);
-			goto err_out;
-		}
-
-		error = of_property_read_u32(np, "reg", &idx);
-		if (error) {
-			dev_err(pmu->dev,
-				"%pOFn: failed to retrieve domain id (reg): %d\n",
-				np, error);
-			goto err_out;
-		}
-		child_domain = pmu->genpd_data.domains[idx];
-
-		error = pm_genpd_add_subdomain(parent_domain, child_domain);
-		if (error) {
-			dev_err(pmu->dev, "%s failed to add subdomain %s: %d\n",
-				parent_domain->name, child_domain->name, error);
-			goto err_out;
-		} else {
-			dev_dbg(pmu->dev, "%s add subdomain: %s\n",
-				parent_domain->name, child_domain->name);
-		}
-
-		rockchip_pm_add_subdomain(pmu, np);
-	}
-
-	return 0;
-
-err_out:
-	of_node_put(np);
-	return error;
-}
-
 static int rockchip_pm_domain_probe(struct platform_device *pdev)
 {
 	struct device *dev = &pdev->dev;
 	struct device_node *np = dev->of_node;
-	struct device_node *node;
 	struct device *parent;
 	struct rockchip_pmu *pmu;
 	const struct of_device_id *match;
 	const struct rockchip_pmu_info *pmu_info;
-	int error;
 
 	if (!np) {
 		dev_err(dev, "device tree node not found\n");

@@ -688,42 +700,9 @@ static int rockchip_pm_domain_probe(struct platform_device *pdev)
 		rockchip_configure_pd_cnt(pmu, pmu_info->gpu_pwrcnt_offset,
 					pmu_info->gpu_power_transition_time);
 
-	error = -ENODEV;
-
-	for_each_available_child_of_node(np, node) {
-		error = rockchip_pm_add_one_domain(pmu, node);
-		if (error) {
-			dev_err(dev, "failed to handle node %pOFn: %d\n",
-				node, error);
-			of_node_put(node);
-			goto err_out;
-		}
-
-		error = rockchip_pm_add_subdomain(pmu, node);
-		if (error < 0) {
-			dev_err(dev, "failed to handle subdomain node %pOFn: %d\n",
-				node, error);
-			of_node_put(node);
-			goto err_out;
-		}
-	}
-
-	if (error) {
-		dev_dbg(dev, "no power domains defined\n");
-		goto err_out;
-	}
-
-	error = of_genpd_add_provider_onecell(np, &pmu->genpd_data);
-	if (error) {
-		dev_err(dev, "failed to add provider: %d\n", error);
-		goto err_out;
-	}
+	rockchip_pm_add_domains(pmu, np, NULL);
 
 	return 0;
-
-err_out:
-	rockchip_pm_domain_cleanup(pmu);
-	return error;
 }
 
 static const struct rockchip_domain_info px30_pm_domains[] = {